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.