heat.linalg.solver

Collection of solvers for systems of linear equations.

Module Contents

cg(A: heat.core.dndarray.DNDarray, b: heat.core.dndarray.DNDarray, x0: heat.core.dndarray.DNDarray, out: heat.core.dndarray.DNDarray | None = None) → heat.core.dndarray.DNDarray[source]

Conjugate gradients method for solving a system of linear equations :math: Ax = b

Parameters:

  • A (DNDarray) – 2D symmetric, positive definite Matrix

  • b (DNDarray) – 1D vector

  • x0 (DNDarray) – Arbitrary 1D starting vector

  • out (DNDarray, optional) – Output Vector

lanczos(A: heat.core.dndarray.DNDarray, m: int, v0: heat.core.dndarray.DNDarray | None = None, V_out: heat.core.dndarray.DNDarray | None = None, T_out: heat.core.dndarray.DNDarray | None = None) → Tuple[heat.core.dndarray.DNDarray, heat.core.dndarray.DNDarray][source]

The Lanczos algorithm is an iterative approximation of the solution to the eigenvalue problem, as an adaptation of power methods to find the m “most useful” (tending towards extreme highest/lowest) eigenvalues and eigenvectors of an \(n \times n\) Hermitian matrix, where often \(m<<n\). It returns two matrices \(V\) and \(T\), where:

  • \(V\) is a Matrix of size \(n\times m\), with orthonormal columns, that span the Krylow subspace n

  • \(T\) is a Tridiagonal matrix of size \(m\times m\), with coefficients \(\alpha_1,..., \alpha_n\) on the diagonal and coefficients \(\beta_1,...,\beta_{n-1}\) on the side-diagonalsn

Parameters:

  • A (DNDarray) – 2D Hermitian (if complex) or symmetric positive-definite matrix. Only distribution along axis 0 is supported, i.e. A.split must be 0 or None.

  • m (int) – Number of Lanczos iterations

  • v0 (DNDarray, optional) – 1D starting vector of Euclidean norm 1. If not provided, a random vector will be used to start the algorithm

  • V_out (DNDarray, optional) – Output Matrix for the Krylow vectors, Shape = (n, m), dtype=A.dtype, must be initialized to zero

  • T_out (DNDarray, optional) – Output Matrix for the Tridiagonal matrix, Shape = (m, m), must be initialized to zero

solve(A: heat.core.dndarray.DNDarray, b: heat.core.dndarray.DNDarray, out: heat.core.dndarray.DNDarray | None = None) → heat.core.dndarray.DNDarray[source]

Computes the solution of a square system of linear equations with a unique solution.

Letting \(\mathbb{K}\) be \(\mathbb{R}\) or \(\mathbb{C}\), this function computes the solution \(X \in \mathbb{K}^{n \times k}\) of the linear system associated to \(A \in \mathbb{K}^{n \times n}, B \in \mathbb{K}^{n \times k}\), which is defined as

\[AX = B\]

Supports inputs of integer, float, double, cfloat and cdouble dtypes. Also supports batches of matrices, and if the inputs are batches of matrices then the output has the same batch dimensions.

Letting * be zero or more batch dimensions,

  • If A has shape (*, n, n) and B has shape (*, n) (a batch of vectors) or shape (*, n, k) (a batch of matrices or “multiple right-hand sides”), this function returns X of shape (*, n) or (*, n, k) respectively.

  • Otherwise, if A has shape (*, n, n) and B has shape (n,) or (n, k), B is broadcast to have shape (*, n) or (*, n, k) respectively. This function then returns the solution of the resulting batch of systems of linear equations.

Note

A and b may only be distributed in the batch dimensions. If both are split, they must be split along matching batch axes.

See also

torch.linalg.solve() is called under the hood on the local data. This docstring is also heavily inspired by the docstring of this function.

Parameters:

  • A (DNDarray) – Matrix to be inverted of shape (*, n, n) where * is zero or more batch dimensions

  • b (DNDarray) – Right-hand side of shape (*, n) or (*, n, k) or (n,) or (n, k)

  • out (DNDarray, optional) – Output Vector

  • Examples:: – 

    >>> A = ht.random.randn(3, 3)
>>> b = ht.random.randn(3)
>>> x = ht.linalg.solve(A, b)
>>> ht.allclose(A @ x, b, atol=1e-5)
True
>>> A = ht.random.randn(2, 3, 3, split=0)
>>> B = ht.random.randn(2, 3, 4, split=0)
>>> X = ht.linalg.solve(A, B)
>>> X.shape
(2, 3, 4)
>>> ht.allclose(A @ X, B, atol=1e-5)
True
>>> A = ht.random.randn(2, 3, 3, split=None)
>>> B = ht.random.randn(2, 3, 4, split=2)
>>> X = ht.linalg.solve(A, B)
>>> X.split
2
>>> ht.allclose(A @ X, B, atol=1e-5)
True

    >>> A = ht.random.randn(2, 3, 3, split=0)
>>> b = ht.random.randn(3, 1)
>>> x = ht.linalg.solve(A, b) # b is broadcast to size (2, 3, 1)
>>> x.shape
(2, 3, 1)
>>> x.split
0
>>> ht.allclose((A @ x).resplit_(None), b, atol=1e-5)
True

solve_triangular(A: heat.core.dndarray.DNDarray, b: heat.core.dndarray.DNDarray) → heat.core.dndarray.DNDarray[source]

Solver for (possibly batched) upper triangular systems of linear equations: it returns x in Ax = b, where A is a (possibly batched) upper triangular matrix and b a (possibly batched) vector or matrix of suitable shape, both provided as input to the function. The implementation builts on the corresponding solver in PyTorch and implements an memory-distributed, MPI-parallel block-wise version thereof.

Parameters:

  • A (DNDarray) – An upper triangular invertible square (n x n) matrix or a batch thereof, i.e. a DNDarray of shape (…, n, n).

  • b (DNDarray) – a (possibly batched) n x k matrix, i.e. an DNDarray of shape (…, n, k), where the batch-dimensions denoted by … need to coincide with those of A. (Batched) Vectors have to be provided as … x n x 1 matrices and the split dimension of b must the second last dimension if not None.

Note

Since such a check might be computationally expensive, we do not check whether A is indeed upper triangular. If you require such a check, please open an issue on our GitHub page and request this feature.