Skip to content

matfree.lanczos

matfree.lanczos

All things Lanczos' algorithm.

This includes stochastic Lanczos quadrature (extending the integrands in hutchinson to those that implement stochastic Lanczos quadrature), Lanczos-implementations of matrix-function-vector products, and various Lanczos-decompositions of matrices.

Examples:

>>> import jax.random
>>> import jax.numpy as jnp
>>>
>>> jnp.set_printoptions(1)
>>>
>>> M = jax.random.normal(jax.random.PRNGKey(1), shape=(10, 10))
>>> A = M.T @ M
>>> v = jax.random.normal(jax.random.PRNGKey(2), shape=(10,))
>>>
>>> # Compute a matrix-logarithm with Lanczos' algorithm
>>> matfun_vec = funm_vector_product_spd(jnp.log, 4, lambda s: A @ s)
>>> matfun_vec(v)
Array([-4. , -2.1, -2.7, -1.9, -1.3, -3.5, -0.5, -0.1,  0.3,  1.5],      dtype=float32)

matfree.lanczos.alg_bidiag_full_reortho(Av: Callable, vA: Callable, depth, /, matrix_shape, validate_unit_2_norm=False)

Construct an implementation of bidiagonalisation.

Uses pre-allocation. Fully reorthogonalise vectors at every step.

Works for arbitrary matrices. No symmetry required.

Decompose a matrix into a product of orthogonal-bidiagonal-orthogonal matrices. Use this algorithm for approximate singular value decompositions.

matfree.lanczos.alg_tridiag_full_reortho(Av: Callable, depth, /, validate_unit_2_norm=False) -> Callable

Construct an implementation of tridiagonalisation.

Uses pre-allocation. Fully reorthogonalise vectors at every step.

This algorithm assumes a symmetric matrix.

Decompose a matrix into a product of orthogonal-tridiagonal-orthogonal matrices. Use this algorithm for approximate eigenvalue decompositions.

matfree.lanczos.funm_vector_product_spd(matfun, order, matvec)

Implement a matrix-function-vector product via Lanczos' algorithm.

This algorithm uses Lanczos' tridiagonalisation with full re-orthogonalisation and therefore applies only to symmetric, positive definite matrices.

matfree.lanczos.integrand_product(matfun, depth, matvec, vecmat)

Construct the integrand for the trace of a function of a matrix-product.

Instead of the trace of a function of a matrix, compute the trace of a function of the product of matrices. Here, "product" refers to \(X = A^\top A\).

matfree.lanczos.integrand_product_logdet(depth, matvec, vecmat)

Construct the integrand for the log-determinant of a matrix-product.

Here, "product" refers to \(X = A^\top A\).

matfree.lanczos.integrand_product_schatten_norm(power, depth, matvec, vecmat)

Construct the integrand for the p-th power of the Schatten-p norm.

matfree.lanczos.integrand_spd(matfun, order, matvec)

Quadratic form for stochastic Lanczos quadrature.

This function assumes a symmetric, positive definite matrix.

matfree.lanczos.integrand_spd_logdet(order, matvec)

Construct the integrand for the log-determinant.

This function assumes a symmetric, positive definite matrix.

matfree.lanczos.svd_approx(v0: Array, depth: int, Av: Callable, vA: Callable, matrix_shape: Tuple[int, ...])

Approximate singular value decomposition.

Uses GKL with full reorthogonalisation to bi-diagonalise the target matrix and computes the full SVD of the (small) bidiagonal matrix.

Parameters:

  • v0 (Array) –

    Initial vector for Golub-Kahan-Lanczos bidiagonalisation.

  • depth (int) –

    Depth of the Krylov space constructed by Golub-Kahan-Lanczos bidiagonalisation. Choosing depth = min(nrows, ncols) - 1 would yield behaviour similar to e.g. np.linalg.svd.

  • Av (Callable) –

    Matrix-vector product function.

  • vA (Callable) –

    Vector-matrix product function.

  • matrix_shape (Tuple[int, ...]) –

    Shape of the matrix involved in matrix-vector and vector-matrix products.