:mod:`heat.linalg.solver`
==============================
.. py:module:: heat.core.linalg.solver

.. autoapi-nested-parse::

   Collection of solvers for systems of linear equations.



Module Contents
---------------


.. function:: cg(A: heat.core.dndarray.DNDarray, b: heat.core.dndarray.DNDarray, x0: heat.core.dndarray.DNDarray, out: Optional[heat.core.dndarray.DNDarray] = None) -> heat.core.dndarray.DNDarray

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

   :param A: 2D symmetric, positive definite Matrix
   :type A: DNDarray
   :param b: 1D vector
   :type b: DNDarray
   :param x0: Arbitrary 1D starting vector
   :type x0: DNDarray
   :param out: Output Vector
   :type out: DNDarray, optional


.. function:: lanczos(A: heat.core.dndarray.DNDarray, m: int, v0: Optional[heat.core.dndarray.DNDarray] = None, V_out: Optional[heat.core.dndarray.DNDarray] = None, T_out: Optional[heat.core.dndarray.DNDarray] = None) -> Tuple[heat.core.dndarray.DNDarray, heat.core.dndarray.DNDarray]

   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 :math:`n \times n` Hermitian matrix, where often :math:`m<<n`.
   It returns two matrices :math:`V` and :math:`T`, where:

       - :math:`V` is a Matrix of size :math:`n\times m`, with orthonormal columns, that span the Krylow subspace \n
       - :math:`T` is a Tridiagonal matrix of size :math:`m\times m`, with coefficients :math:`\alpha_1,..., \alpha_n`
         on the diagonal and coefficients :math:`\beta_1,...,\beta_{n-1}` on the side-diagonals\n

   :param A: 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`.
   :type A: DNDarray
   :param m: Number of Lanczos iterations
   :type m: int
   :param v0: 1D starting vector of Euclidean norm 1. If not provided, a random vector will be used to start the algorithm
   :type v0: DNDarray, optional
   :param V_out: Output Matrix for the Krylow vectors, Shape = (n, m), dtype=A.dtype, must be initialized to zero
   :type V_out: DNDarray, optional
   :param T_out: Output Matrix for the Tridiagonal matrix, Shape = (m, m), must be initialized to zero
   :type T_out: DNDarray, optional


.. function:: solve(A: heat.core.dndarray.DNDarray, b: heat.core.dndarray.DNDarray, out: Optional[heat.core.dndarray.DNDarray] = None) -> heat.core.dndarray.DNDarray

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

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

   .. math:: 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 :attr:`A` has shape `(*, n, n)` and :attr:`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 :attr:`A` has shape `(*, n, n)` and  :attr:`B` has shape `(n,)`  or `(n, k)`, :attr:`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.

   .. seealso::

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


   :param A: Matrix to be inverted of shape `(*, n, n)` where `*` is zero or more batch dimensions
   :type A: DNDarray
   :param b: Right-hand side of shape `(*, n)` or  `(*, n, k)` or `(n,)` or `(n, k)`
   :type b: DNDarray
   :param out: Output Vector
   :type out: DNDarray, optional
   :param 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


.. function:: solve_triangular(A: heat.core.dndarray.DNDarray, b: heat.core.dndarray.DNDarray) -> heat.core.dndarray.DNDarray

   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.

   :param A: An upper triangular invertible square (n x n) matrix or a batch thereof,  i.e. a ``DNDarray`` of shape `(..., n, n)`.
   :type A: DNDarray
   :param b: 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.
   :type b: DNDarray

   .. 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.