Linear Algebra

Inverses & Solves

cola.linalg.inv(A, alg=Auto())[source]

(lazily) computes the inverse of a linear operator, equivalent to solve.

Parameters:

  • A (LinearOperator) – The linear operator to compute the inverse of.

  • alg (Algorithm, optional) – The algorithm to use for the solves.

Returns:

The inverse of the linear operator.

Return type:

LinearOperator

Example

>>> A = MyLinearOperator()
>>> x = cola.inverse(A, alg=CG(tol=1e-3)) @ b
cola.linalg.solve(A, b, alg=Auto())[source]

Computes Linear solve of a linear operator. Equivalent to cola.inv

Parameters:

  • A (LinearOperator) – The linear operator to compute the inverse of.

  • b (Array) – The right hand side of the linear system of shape (d, ) or (d, k)

  • alg (Algorithm, optional) – The algorithm to use for the solves.

Returns:

The solution of the linear system of shape (d, ) or (d, k)

Return type:

Array

Example

>>> A = MyLinearOperator()
>>> x = cola.solve(A, b, alg=Auto(max_iters=10, pbar=True))

Algorithms

class cola.linalg.Auto[source]

Bases: SimpleNamespace, Algorithm

  • if A is PSD and small, uses Cholesky

  • if A is not PSD and small, uses LU

  • if A is PSD and large, uses CG

  • if A is not PSD and large, uses GMRES

class cola.linalg.CG(tol=1e-06, max_iters=1000, pbar=False, x0=None, P=None)[source]

Bases: Algorithm

Conjugate gradients algorithm Solves Ax=b or AX=B (multiple rhs).

The runtime is bounded by \(O(\sqrt{\kappa})\) and it uses \(O(n)\) memory. Where \(\kappa\) is the condition number of the linear operator. Optionally, you can use a preconditioner (approx of A⁻¹) to accelerate convergence.

Parameters:

  • tol (float, optional) – Relative error tolerance.

  • max_iters (int, optional) – The maximum number of iterations to run.

  • pbar (bool, optional) – Whether to show progress bar.

  • x0 (Array, optional) – (n,) or (n, b) guess for initial solution.

  • P (LinearOperator, optional) – Preconditioner. Defaults to the identity.

class cola.linalg.GMRES(tol=1e-06, max_iters=1000, pbar=False, x0=None, P=None)[source]

Bases: Algorithm

Generalized Minimal Residual algorith (GMRES) for soving Ax=b or AX=B (multiple rhs).

The runtime is bounded by \(O(\sqrt{\kappa})\) and it uses \(O(m n)\) memory. Where \(\kappa\) is the condition number of the linear operator, n is the size of A and m represents the max iters. Optionally, you can use a preconditioner (approx of A⁻¹) to accelerate convergence.

Parameters:

  • tol (float, optional) – Relative error tolerance for Arnoldi.

  • max_iters (int, optional) – The maximum number of iterations to run in Arnoldi.

  • pbar (bool, optional) – Whether to show progress bar.

  • x0 (Array, optional) – (n,) or (n, b) guess for initial solution.

  • P (LinearOperator, optional) – Preconditioner. Defaults to the identity.

class cola.linalg.LU[source]

Bases: Algorithm

LU algorithm for decomposing a general square operator as \(A = PLU\), where \(P\) is a permutation operator, \(L\) is a lower triangular operator and \(U\) is an upper triangular operator.

Example

>>> A = MyLinearOperator()
>>> P,L,U = cola.linalg.decompositions.LU()(A)
class cola.linalg.Cholesky[source]

Bases: Algorithm

Cholesky algorithm for decomposing a positive definite operator as \(A = L L^{*}\), where \(L\) is a lower triangular operator.

Example

>>> A = MyLinearOperator()
>>> L = cola.linalg.decompositions.Cholesky()(A)

Eigenvalues and Eigenvectors

cola.linalg.eig(A, k, which='LM', alg=Auto())[source]

Computes eigenvalues and eigenvectors of a linear operator.

Parameters:

  • A (LinearOperator) – The linear operator for which eigenvalues and eigenvectors are computed.

  • k (int) – The desired number of eigenvalues and eigenvectors. Must be specified.

  • which (str) – From what part of the spectrum would de eigenvalues be fetched. Default is β€˜LM’ (largest in magnitude) but alternatively you can use β€˜SM’ (smallest in magnitude).

  • alg (Algorithm) – (Auto, Eig, Eigh, Arnoldi, Lanczos)

Returns:

A tuple containing eigenvalues and eigenvectors.

The eigenvalues are given by eig_vals and the eigenvectors are given by eig_vecs.

Return type:

Tuple[Array, LinearOperator]

Example

>>> A = MyLinearOperator()
>>> eig_vals, eig_vecs = eig(A, k=6, which='LM', alg=Auto(tol=1e-4))

Algorithms

class cola.linalg.Auto[source]

Bases: SimpleNamespace, Algorithm

  • if A is Hermitian and small, uses Eigh

  • if A is not Hermitian and small, use Eig

  • if A is Hermitian and large, uses Lanczos

  • if A is not Hermitian and large, uses Arnoldi

class cola.linalg.Arnoldi(start_vector=None, max_iters=1000, tol=1e-06, pbar=False, key=None)[source]

Bases: Algorithm

Arnoldi decomposition for a general square operator, \(H \approx Q^{*} A Q\) where \(H\) is an upper Hessenberg operator. This algorithm is used to approximate eig(A) through eig(H).

Parameters:

  • start_vector (Array, optional) – (n,) or (n, b) vector to start the algorithm.

  • max_iters (int, optional) – The maximum number of iterations to run.

  • tol (float, optional) – Relative error tolerance.

  • pbar (bool, optional) – Whether to show progress bar.

  • key (PRNGKey, optional) – Random key to use for the algorithm. PRNGKey for jax, long integer for numpy or pytorch.

Example

>>> A = MyLinearOperator()
>>> Q,H,info = Arnoldi(max_iters=100,pbar=True)(A)
class cola.linalg.Lanczos(start_vector=None, max_iters=1000, tol=1e-06, pbar=False, key=None)[source]

Bases: Algorithm

Lanczos decomposition for Symmetric (or Hermitian) operators, \(T = Q^{*} A Q\) where \(T\) is a tridiagional operator. This algorithm is used to approximate eig(A) through eig(T).

Parameters:

  • start_vector (Array, optional) – (n,) or (n, b) vector to start the algorithm.

  • max_iters (int, optional) – The maximum number of iterations to run.

  • tol (float, optional) – Relative error tolerance.

  • pbar (bool, optional) – Whether to show progress bar.

  • key (PRNGKey, optional) – Random key to use for the algorithm. PRNGKey for jax, long integer for numpy or pytorch.

Example

>>> A = MyLinearOperator()
>>> Q,T,info = Lanczos(max_iters=100,pbar=True)(A)
class cola.linalg.Eig[source]

Bases: Algorithm

Uses a dense eigendecomposition for a general square operator.

class cola.linalg.Eigh[source]

Bases: Algorithm

Uses a dense eigendecomposition for a real symmetric or complex Hermitian operator.

class cola.linalg.PowerIteration(tol=1e-06, max_iter=100, pbar=False, key=None)[source]

Bases: Algorithm

Simple power iteration algorithm for finding the largest eigenvalue and eigenvector.

Parameters:

  • tol (float, optional) – Relative error tolerance.

  • max_iters (int, optional) – The maximum number of iterations to run.

  • pbar (bool, optional) – Whether to show progress bar.

  • key (PRNGKey, optional) – Random key for reproducibility.

Example

>>> A = MyLinearOperator()
>>> v, eigmax, info = PowerIteration(tol=1e-3)(A)

Matrix Functions: exp, log, sqrt, isqrt, etc.

cola.linalg.exp(A, alg=Auto())[source]

Computes the matrix exponential \(\exp(A)\) of the operator \(A\).

Parameters:

  • f (Callable) – The function to apply.

  • A (LinearOperator) – The linear operator to compute f(A) with.

  • alg (Algorithm) – The algorithm to use (Auto, Eig, Eigh, Lanczos, Arnoldi).

Returns:

a lazy instantiation of \(\exp(A)\)

Return type:

LinearOperator

cola.linalg.exp(A: cola.ops.operators.KronSum, alg: cola.linalg.algorithm_base.Algorithm)[source]
cola.linalg.log(A, alg=Auto())[source]

Computes the matrix logarithm \(log(A)\) of positive definite operator \(A\).

Parameters:

  • A (LinearOperator) – The linear operator to compute f(A) with.

  • alg (Algorithm) – The algorithm to use (Auto, Eig, Eigh, Lanczos, Arnoldi).

Returns:

a lazy instantiation of log(A)

Return type:

LinearOperator

cola.linalg.sqrt(A, alg=Auto())[source]

Computes the square root, \(A^{1/2}\) of an operator \(A\) using the principal branch.

Parameters:

  • A (LinearOperator) – The linear operator to compute f(A) with.

  • alg (Algorithm) – The algorithm to use (Auto, Eig, Eigh, Lanczos, Arnoldi).

Returns:

a lazy instantiation of \(A^{1/2}\)

Return type:

LinearOperator

cola.linalg.isqrt(A, alg=Auto())[source]

Computes the matrix inverse \(A^{-1/2}\) of an operator \(A\) using the principal branch.

Parameters:

  • A (LinearOperator) – The linear operator to compute f(A) with.

  • alg (Algorithm) – The algorithm to use (Auto, Eig, Eigh, Lanczos, Arnoldi).

Returns:

a lazy instantiation of \(A^{-1/2}\)

Return type:

LinearOperator

cola.linalg.pow(A, alpha, alg=Auto())[source]

Computes the matrix power \(A^{\alpha}\) of an operator \(A\), where \(\alpha\) is the coefficient.

Parameters:

  • A (LinearOperator) – The linear operator to compute f(A) with.

  • alg (Algorithm) – The algorithm to use (Auto, Eig, Eigh, Lanczos, Arnoldi).

Returns:

a lazy instantiation of \(A^{\alpha}\)

Return type:

LinearOperator

cola.linalg.pow(A: cola.ops.operators.Kronecker, alpha: numbers.Number, alg: cola.linalg.algorithm_base.Algorithm)[source]
cola.linalg.apply_unary(f, A, alg=Auto())[source]

Generic apply a unary function \(f\) to a linear operator \(A\). That is, \(f(A)\), defined through the taylor expansion:

\(f(A) = \sum_{k=0}^\infty \frac{f^{(k)}(0)}{k!}A^k\).

Parameters:

  • f (Callable) – The function to apply.

  • A (LinearOperator) – The linear operator to compute f(A) with.

  • alg (Algorithm) – The algorithm to use (Auto, Eig, Eigh, Lanczos, Arnoldi).

Returns:

a lazy instantiation of \(f(A)\)

Return type:

LinearOperator

Algorithms

class cola.linalg.Auto[source]

Bases: SimpleNamespace, Algorithm

  • if A is Hermitian and small, uses Eigh

  • if A is not Hermitian and small, uses Eig

  • if A is Hermitian and large, uses Lanczos

  • if A is not Hermitian and large, uses Arnoldi

class cola.linalg.Lanczos(start_vector=None, max_iters=1000, tol=1e-06, pbar=False, key=None)[source]

Bases: Algorithm

Lanczos decomposition for Symmetric (or Hermitian) operators, \(T = Q^{*} A Q\) where \(T\) is a tridiagional operator. This algorithm is used to approximate eig(A) through eig(T).

Parameters:

  • start_vector (Array, optional) – (n,) or (n, b) vector to start the algorithm.

  • max_iters (int, optional) – The maximum number of iterations to run.

  • tol (float, optional) – Relative error tolerance.

  • pbar (bool, optional) – Whether to show progress bar.

  • key (PRNGKey, optional) – Random key to use for the algorithm. PRNGKey for jax, long integer for numpy or pytorch.

Example

>>> A = MyLinearOperator()
>>> Q,T,info = Lanczos(max_iters=100,pbar=True)(A)
class cola.linalg.Arnoldi(start_vector=None, max_iters=1000, tol=1e-06, pbar=False, key=None)[source]

Bases: Algorithm

Arnoldi decomposition for a general square operator, \(H \approx Q^{*} A Q\) where \(H\) is an upper Hessenberg operator. This algorithm is used to approximate eig(A) through eig(H).

Parameters:

  • start_vector (Array, optional) – (n,) or (n, b) vector to start the algorithm.

  • max_iters (int, optional) – The maximum number of iterations to run.

  • tol (float, optional) – Relative error tolerance.

  • pbar (bool, optional) – Whether to show progress bar.

  • key (PRNGKey, optional) – Random key to use for the algorithm. PRNGKey for jax, long integer for numpy or pytorch.

Example

>>> A = MyLinearOperator()
>>> Q,H,info = Arnoldi(max_iters=100,pbar=True)(A)
class cola.linalg.Eig[source]

Bases: Algorithm

Uses a dense eigendecomposition for a general square operator.

class cola.linalg.Eigh[source]

Bases: Algorithm

Uses a dense eigendecomposition for a real symmetric or complex Hermitian operator.

Trace, Diagonal, Frobenius Norm

cola.linalg.trace(A, alg=Auto())[source]

Compute the trace of a linear operator tr(A).

Can use either the \(O(\tfrac{1}{\delta^2})\) time stochastic estimation (alg=Hutch()) or a deterministic \(O(n)\) time algorithm (alg =Exact()).

If only unbiased estimates of the diagonal are needed, use the Hutchinson algorithm.

Parameters:

  • A (LinearOperator) – The linear operator to compute the diagonal of.

  • alg (Algorithm, optional) – The algorithm to use for the diagonal computation (Hutch or Exact).

Returns:

trace

Return type:

Array

cola.linalg.diag(A, k=0, alg=Auto())[source]

Extract the (kth) diagonal of a linear operator.

Can use either the \(O(\tfrac{1}{\delta^2})\) time stochastic estimation (alg=Hutch()) or a deterministic \(O(n)\) time algorithm (alg =Exact()).

If only unbiased estimates of the diagonal are needed, use the Hutchinson algorithm.

Parameters:

  • A (LinearOperator) – The linear operator to compute the diagonal of.

  • k (int, optional) – Specify to compute the kth off diagonal diagonal.

  • alg (Algorithm, optional) – The algorithm to use for the diagonal computation (Hutch or Exact).

Returns:

diag

Return type:

Array

Algorithms

class cola.linalg.Auto[source]

Bases: SimpleNamespace, Algorithm

  • if \(\tfrac{1}{\sqrt{10n}} < \epsilon\) use Hutch \(O(\tfrac{1}{\delta^2})\)

  • otherwise use Exact \(O(n)\)

class cola.linalg.Hutch(tol=0.03, max_iters=10000, bs=100, rand='normal', pbar=False, key=None)[source]

Hutchinson’s algorithm for estimating the trace of a matrix function. Basically, this algorithm uses random probes to approximate \(\text{tr}(f(A))\).

Parameters:

  • tol (float, optional) – Approximation relative tolerance.

  • max_iters (int, optional) – The maximum number of iterations to run.

  • bs (int, optional) – Number of probes.

  • rand (str, optional) – type of random probes (either Normal or Rademacher)

  • pbar (bool, optional) – Whether to show progress bar.

  • key (xnp.PRNGKey, optional) – Random key (default None).

class cola.linalg.Exact(bs=100, pbar=False)[source]

Exact algorithm to compute, element-by-element, the diagonal of the matrix. That is, each element of the diagoinal \(A_{i,i}\) is computed as \(e_{i}^{T} A(e_{i})\). For efficiency, this procedure is done through blocks of elements, \(A [e_{i_{1}}, \cdots, e_{i_{B}}]\) where \(B\) is the block size.

Parameters:

  • bs (int, optional) – Block size.

  • pbar (bool, optional) – Whether to show progress bar.

class cola.linalg.HutchPP(tol=0.03, max_iters=10000, bs=100, rand='normal', pbar=False, key=None)[source]

Hutch++ is an improvement on the Hutchinson’s estimator introduced in Meyer et al. 2020: Hutch++: Optimal Stochastic Trace Estimator.

Parameters:

  • tol (float, optional) – Approximation relative tolerance.

  • max_iters (int, optional) – The maximum number of iterations to run.

  • bs (int, optional) – Number of probes.

  • rand (str, optional) – type of random probes (either Normal or Rademacher)

  • pbar (bool, optional) – Whether to show progress bar.

  • key (xnp.PRNGKey, optional) – Random key (default None).

Log Determinants

cola.linalg.logdet(A, log_alg=Auto(), trace_alg=Auto())[source]

Computes log determinant of a linear operator.

For large inputs (or with method=’iterative’), uses either \(O(\tfrac{1}{\delta^2}\log(1/\epsilon))\) time stochastic algorithm (SLQ) where \(\epsilon=\) tol is the bias and \(\delta=\) vtol is the standard deviation of the estimate, or a deterministic \(O(n\log(1/\epsilon))\) time algorithm if \(\delta < 1/\sqrt{10n}\).

Parameters:

  • A (LinearOperator) – The linear operator to compute the logdet of.

  • log_alg (Algorithm, optional) – The algorithm to use for the log. Specify LU() or Cholesky() for a dense approach.

  • trace_alg (Algorithm, optional) – The algorithm to use for the trace computation.

Returns:

logdet

Return type:

Array

cola.linalg.slogdet(A, log_alg=Auto(), trace_alg=Auto())[source]

Computes sign and logdet of a linear operator. such that det(A) = sign(A) exp(logdet(A))

Parameters:

  • A (LinearOperator) – The linear operator to compute the logdet of.

  • log_alg (Algorithm, optional) – The algorithm to use for the log. Specify LU() or Cholesky() for a dense approach.

  • trace_alg (Algorithm, optional) – The algorithm to use for the trace computation.

Returns:

sign, logdet: logdet

Return type:

Tuple[Array, Array]

#

Log Algs

class cola.linalg.Auto[source]

Bases: SimpleNamespace, Algorithm

  • if A is PSD and small, uses Cholesky

  • if A is not PSD and small, uses LU

  • if A is PSD and large, uses Lanczos

  • if A is not PSD and large, uses Arnoldi

class cola.linalg.Lanczos(start_vector=None, max_iters=1000, tol=1e-06, pbar=False, key=None)[source]

Bases: Algorithm

Lanczos decomposition for Symmetric (or Hermitian) operators, \(T = Q^{*} A Q\) where \(T\) is a tridiagional operator. This algorithm is used to approximate eig(A) through eig(T).

Parameters:

  • start_vector (Array, optional) – (n,) or (n, b) vector to start the algorithm.

  • max_iters (int, optional) – The maximum number of iterations to run.

  • tol (float, optional) – Relative error tolerance.

  • pbar (bool, optional) – Whether to show progress bar.

  • key (PRNGKey, optional) – Random key to use for the algorithm. PRNGKey for jax, long integer for numpy or pytorch.

Example

>>> A = MyLinearOperator()
>>> Q,T,info = Lanczos(max_iters=100,pbar=True)(A)
class cola.linalg.Arnoldi(start_vector=None, max_iters=1000, tol=1e-06, pbar=False, key=None)[source]

Bases: Algorithm

Arnoldi decomposition for a general square operator, \(H \approx Q^{*} A Q\) where \(H\) is an upper Hessenberg operator. This algorithm is used to approximate eig(A) through eig(H).

Parameters:

  • start_vector (Array, optional) – (n,) or (n, b) vector to start the algorithm.

  • max_iters (int, optional) – The maximum number of iterations to run.

  • tol (float, optional) – Relative error tolerance.

  • pbar (bool, optional) – Whether to show progress bar.

  • key (PRNGKey, optional) – Random key to use for the algorithm. PRNGKey for jax, long integer for numpy or pytorch.

Example

>>> A = MyLinearOperator()
>>> Q,H,info = Arnoldi(max_iters=100,pbar=True)(A)
class cola.linalg.Cholesky[source]

Bases: Algorithm

Cholesky algorithm for decomposing a positive definite operator as \(A = L L^{*}\), where \(L\) is a lower triangular operator.

Example

>>> A = MyLinearOperator()
>>> L = cola.linalg.decompositions.Cholesky()(A)
class cola.linalg.LU[source]

Bases: Algorithm

LU algorithm for decomposing a general square operator as \(A = PLU\), where \(P\) is a permutation operator, \(L\) is a lower triangular operator and \(U\) is an upper triangular operator.

Example

>>> A = MyLinearOperator()
>>> P,L,U = cola.linalg.decompositions.LU()(A)

Trace Algs

class cola.linalg.Auto[source]

  • if \(\tfrac{1}{\sqrt{10n}} < \epsilon\) use Hutch \(O(\tfrac{1}{\delta^2})\)

  • otherwise use Exact \(O(n)\)

class cola.linalg.Hutch(tol=0.03, max_iters=10000, bs=100, rand='normal', pbar=False, key=None)[source]

Hutchinson’s algorithm for estimating the trace of a matrix function. Basically, this algorithm uses random probes to approximate \(\text{tr}(f(A))\).

Parameters:

  • tol (float, optional) – Approximation relative tolerance.

  • max_iters (int, optional) – The maximum number of iterations to run.

  • bs (int, optional) – Number of probes.

  • rand (str, optional) – type of random probes (either Normal or Rademacher)

  • pbar (bool, optional) – Whether to show progress bar.

  • key (xnp.PRNGKey, optional) – Random key (default None).

class cola.linalg.Exact(bs=100, pbar=False)[source]

Exact algorithm to compute, element-by-element, the diagonal of the matrix. That is, each element of the diagoinal \(A_{i,i}\) is computed as \(e_{i}^{T} A(e_{i})\). For efficiency, this procedure is done through blocks of elements, \(A [e_{i_{1}}, \cdots, e_{i_{B}}]\) where \(B\) is the block size.

Parameters:

  • bs (int, optional) – Block size.

  • pbar (bool, optional) – Whether to show progress bar.

class cola.linalg.HutchPP(tol=0.03, max_iters=10000, bs=100, rand='normal', pbar=False, key=None)[source]

Hutch++ is an improvement on the Hutchinson’s estimator introduced in Meyer et al. 2020: Hutch++: Optimal Stochastic Trace Estimator.

Parameters:

  • tol (float, optional) – Approximation relative tolerance.

  • max_iters (int, optional) – The maximum number of iterations to run.

  • bs (int, optional) – Number of probes.

  • rand (str, optional) – type of random probes (either Normal or Rademacher)

  • pbar (bool, optional) – Whether to show progress bar.

  • key (xnp.PRNGKey, optional) – Random key (default None).