Trajectree 0.0.0__py3-none-any.whl → 0.0.1__py3-none-any.whl

trajectree/quimb/quimb/linalg/base_linalg.py

@@ -0,0 +1,719 @@
+"""Backend agnostic functions for solving matrices either fully or partially.
+"""
+import functools
+import warnings
+
+import numpy as np
+import scipy.linalg as sla
+import scipy.sparse.linalg as spla
+
+from ..utils import raise_cant_find_library_function
+from ..core import qarray, dag, issparse, isdense, vdot, ldmul
+from .numpy_linalg import (
+    eig_numpy,
+    eigs_numpy,
+    svds_numpy,
+)
+
+from .scipy_linalg import (
+    eigs_scipy,
+    eigs_lobpcg,
+    eigs_primme,
+    svds_scipy,
+    svds_primme,
+)
+from . import SLEPC4PY_FOUND
+
+if SLEPC4PY_FOUND:
+    from .mpi_launcher import (
+        eigs_slepc_spawn,
+        mfn_multiply_slepc_spawn,
+        svds_slepc_spawn,
+    )
+    from .slepc_linalg import eigs_slepc, svds_slepc, mfn_multiply_slepc
+else:  # pragma: no cover
+    eigs_slepc = raise_cant_find_library_function("slepc4py")
+    eigs_slepc_spawn = raise_cant_find_library_function("slepc4py")
+    svds_slepc = raise_cant_find_library_function("slepc4py")
+    svds_slepc_spawn = raise_cant_find_library_function("slepc4py")
+    mfn_multiply_slepc = raise_cant_find_library_function("slepc4py")
+    mfn_multiply_slepc_spawn = raise_cant_find_library_function("slepc4py")
+
+
+# --------------------------------------------------------------------------- #
+#                          Partial eigendecomposition                         #
+# --------------------------------------------------------------------------- #
+
+
+def choose_backend(A, k, int_eps=False, B=None):
+    """Pick a backend automatically for partial decompositions."""
+    # LinOps -> not possible to simply convert to dense or use MPI processes
+    A_is_linop = isinstance(A, spla.LinearOperator)
+    B_is_linop = isinstance(B, spla.LinearOperator)
+
+    # small array or large part of subspace requested
+    small_d_big_k = A.shape[0] ** 2 / k < (10000 if int_eps else 2000)
+
+    if small_d_big_k and not (A_is_linop or B_is_linop):
+        return "NUMPY"
+
+    # slepc seems faster for sparse, dense and LinearOperators
+    if SLEPC4PY_FOUND and not B_is_linop:
+        # only spool up an mpi pool for big sparse matrices though
+        if issparse(A) and A.nnz > 10000:
+            return "SLEPC"
+
+        return "SLEPC-NOMPI"
+
+    return "SCIPY"
+
+
+_EIGS_METHODS = {
+    "NUMPY": eigs_numpy,
+    "SCIPY": eigs_scipy,
+    "PRIMME": eigs_primme,
+    "LOBPCG": eigs_lobpcg,
+    "SLEPC": eigs_slepc_spawn,
+    "SLEPC-NOMPI": eigs_slepc,
+}
+
+
+def eigensystem_partial(
+    A,
+    k,
+    isherm,
+    *,
+    B=None,
+    which=None,
+    return_vecs=True,
+    sigma=None,
+    ncv=None,
+    tol=None,
+    v0=None,
+    sort=True,
+    backend=None,
+    fallback_to_scipy=False,
+    **backend_opts,
+):
+    """Return a few eigenpairs from an operator.
+
+    Parameters
+    ----------
+    A : sparse, dense or linear operator
+        The operator to solve for.
+    k : int
+        Number of eigenpairs to return.
+    isherm : bool
+        Whether to use hermitian solve or not.
+    B : sparse, dense or linear operator, optional
+        If given, the RHS operator defining a generalized eigen problem.
+    which : {'SA', 'LA', 'LM', 'SM', 'TR'}
+        Where in spectrum to take eigenvalues from (see
+        :func:``scipy.sparse.linalg.eigsh``)
+    return_vecs : bool, optional
+        Whether to return the eigenvectors.
+    sigma : float, optional
+        Which part of spectrum to target, implies which='TR' if which is None.
+    ncv : int, optional
+        number of lanczos vectors, can use to optimise speed
+    tol : None or float
+        Tolerance with which to find eigenvalues.
+    v0 : None or 1D-array like
+        An initial vector guess to iterate with.
+    sort : bool, optional
+        Whether to explicitly sort by ascending eigenvalue order.
+    backend : {'AUTO', 'NUMPY', 'SCIPY',
+               'LOBPCG', 'SLEPC', 'SLEPC-NOMPI'}, optional
+        Which solver to use.
+    fallback_to_scipy : bool, optional
+        If an error occurs and scipy is not being used, try using scipy.
+    backend_opts
+        Supplied to the backend solver.
+
+    Returns
+    -------
+    elk : (k,) array
+        The ``k`` eigenvalues.
+    evk : (d, k) array
+        Array with ``k`` eigenvectors as columns if ``return_vecs``.
+    """
+    settings = {
+        "k": k,
+        "B": B,
+        "which": (
+            "SA"
+            if (which is None) and (sigma is None)
+            else "TR"
+            if (which is None) and (sigma is not None)
+            else which
+        ),
+        "return_vecs": return_vecs,
+        "sigma": sigma,
+        "isherm": isherm,
+        "ncv": ncv,
+        "sort": sort,
+        "tol": tol,
+        "v0": v0,
+    }
+
+    # Choose backend to perform the decompostion
+    bkd = "AUTO" if backend is None else backend.upper()
+    if bkd == "AUTO":
+        bkd = choose_backend(A, k, sigma is not None, B=B)
+
+    try:
+        return _EIGS_METHODS[bkd](A, **settings, **backend_opts)
+
+    # sometimes e.g. lobpcg fails, worth trying scipy
+    except Exception as e:  # pragma: no cover
+        if fallback_to_scipy and (bkd != "SCIPY"):
+            warnings.warn(
+                f"`eigensystem_partial` with backend '{bkd}' failed, trying "
+                "again with scipy. Set ``fallback_to_scipy=False`` to avoid "
+                "this and see the full error."
+            )
+
+            return eigs_scipy(A, **settings, **backend_opts)
+        else:
+            raise e
+
+
+# --------------------------------------------------------------------------- #
+#                        Full eigendecomposition                              #
+# --------------------------------------------------------------------------- #
+
+
+def eigensystem(A, isherm, *, k=-1, sort=True, return_vecs=True, **kwargs):
+    """Find all or some eigenpairs of an operator.
+
+    Parameters
+    ----------
+    A : operator
+        The operator to decompose.
+    isherm : bool
+        Whether the operator is assumed to be hermitian or not.
+    k : int, optional
+        If negative, find all eigenpairs, else perform partial
+        eigendecomposition and find ``k`` pairs. See
+        :func:`~quimb.linalg.base_linalg.eigensystem_partial`.
+    sort : bool, optional
+        Whether to sort the eigenpairs in ascending eigenvalue order.
+    kwargs
+        Supplied to the backend function.
+
+    Returns
+    -------
+    el : (k,) array
+        Eigenvalues.
+    ev : (d, k) array
+        Corresponding eigenvectors as columns of array, such that
+        ``ev @ diag(el) @ ev.H == A``.
+    """
+    if k < 0:
+        return eig_numpy(
+            A, isherm=isherm, sort=sort, return_vecs=return_vecs, **kwargs
+        )
+
+    return eigensystem_partial(
+        A, k=k, isherm=isherm, sort=sort, return_vecs=return_vecs, **kwargs
+    )
+
+
+eig = functools.partial(eigensystem, isherm=False, return_vecs=True)
+eigh = functools.partial(eigensystem, isherm=True, return_vecs=True)
+eigvals = functools.partial(eigensystem, isherm=False, return_vecs=False)
+eigvalsh = functools.partial(eigensystem, isherm=True, return_vecs=False)
+
+
+@functools.wraps(eigensystem)
+def eigenvectors(A, isherm, *, sort=True, **kwargs):
+    return eigensystem(A, isherm=isherm, sort=sort, **kwargs)[1]
+
+
+eigvecs = functools.partial(eigenvectors, isherm=False)
+eigvecsh = functools.partial(eigenvectors, isherm=True)
+
+
+def groundstate(ham, **kwargs):
+    """Alias for finding lowest eigenvector only."""
+    return eigvecsh(ham, k=1, which="SA", **kwargs)
+
+
+def groundenergy(ham, **kwargs):
+    """Alias for finding lowest eigenvalue only."""
+    return eigvalsh(ham, k=1, which="SA", **kwargs)[0]
+
+
+def bound_spectrum(A, backend="auto", **kwargs):
+    """Return the smallest and largest eigenvalue of hermitian operator ``A``."""
+    el_min = eigvalsh(A, k=1, which="SA", backend=backend, **kwargs)[0]
+    el_max = eigvalsh(A, k=1, which="LA", backend=backend, **kwargs)[0]
+    return el_min, el_max
+
+
+def _rel_window_to_abs_window(el_min, el_max, w_0, w_sz=None):
+    """Convert min/max eigenvalues and relative window to absolute values.
+
+    Parameters
+    ----------
+    el_min : float
+        Smallest eigenvalue.
+    el_max : float
+        Largest eigenvalue.
+    w_0 : float [0.0 - 1.0]
+        Relative window centre.
+    w_sz : float, optional
+        Relative window width.
+
+    Returns
+    -------
+    l_0[, l_min, l_max]:
+        Absolute value of centre of window, lower and upper intervals if a
+        window size is specified.
+    """
+    el_range = el_max - el_min
+    el_w_0 = el_min + w_0 * el_range
+    if w_sz is not None:
+        el_w_min = el_w_0 - w_sz * el_range / 2
+        el_w_max = el_w_0 + w_sz * el_range / 2
+        return el_w_0, el_w_min, el_w_max
+    return el_w_0
+
+
+def eigh_window(
+    A,
+    w_0,
+    k,
+    w_sz=None,
+    backend="AUTO",
+    return_vecs=True,
+    offset_const=1 / 104729,
+    **kwargs,
+):
+    """Return mid-spectrum eigenpairs from a hermitian operator.
+
+    Parameters
+    ----------
+    A : (d, d) operator
+        Operator to retrieve eigenpairs from.
+    w_0 : float [0.0, 1.0]
+        Relative window centre to retrieve eigenpairs from.
+    k : int
+        Target number of eigenpairs to retrieve.
+    w_sz : float, optional
+        Relative maximum window width within which to keep eigenpairs.
+    backend : str, optional
+        Which :func:`~quimb.eigh` backend to use.
+    return_vecs : bool, optional
+        Whether to return eigenvectors as well.
+    offset_const : float, optional
+        Small fudge factor (relative to window range) to avoid 1 / 0 issues.
+
+    Returns
+    -------
+    el : (k,) array
+        Eigenvalues around w_0.
+    ev : (d, k) array
+        The eigenvectors, if ``return_vecs=True``.
+    """
+    w_sz = w_sz if w_sz is not None else 1.1
+
+    if isdense(A) or backend.upper() == "NUMPY":
+        if return_vecs:
+            lk, vk = eigh(A.toarray() if issparse(A) else A, **kwargs)
+        else:
+            lk = eigvalsh(A.toarray() if issparse(A) else A, **kwargs)
+
+        lmin, lmax = lk[0], lk[-1]
+        l_w0, l_wmin, l_wmax = _rel_window_to_abs_window(lmin, lmax, w_0, w_sz)
+
+    else:
+        lmin, lmax = bound_spectrum(A, backend=backend, **kwargs)
+        l_w0, l_wmin, l_wmax = _rel_window_to_abs_window(lmin, lmax, w_0, w_sz)
+        l_w0 += (lmax - lmin) * offset_const  # for 1/0 issues
+
+        if return_vecs:
+            lk, vk = eigh(A, k=k, sigma=l_w0, backend=backend, **kwargs)
+        else:
+            lk = eigvalsh(A, k=k, sigma=l_w0, backend=backend, **kwargs)
+
+    # Trim eigenpairs from beyond window
+    in_window = (lk > l_wmin) & (lk < l_wmax)
+
+    if return_vecs:
+        return lk[in_window], vk[:, in_window]
+
+    return lk[in_window]
+
+
+def eigvalsh_window(*args, **kwargs):
+    """Alias for only finding the eigenvalues in a relative window."""
+    return eigh_window(*args, return_vecs=False, **kwargs)
+
+
+def eigvecsh_window(*args, **kwargs):
+    """Alias for only finding the eigenvectors in a relative window."""
+    return eigh_window(*args, return_vecs=True, **kwargs)[1]
+
+
+# -------------------------------------------------------------------------- #
+# Partial singular value decomposition                                       #
+# -------------------------------------------------------------------------- #
+
+
+def svd(A, return_vecs=True):
+    """Compute full singular value decomposition of an operator, using numpy.
+
+    Parameters
+    ----------
+    A : (m, n) array
+        The operator.
+    return_vecs : bool, optional
+        Whether to return the singular vectors.
+
+    Returns
+    -------
+    U : (m, k) array
+        Left singular vectors (if ``return_vecs=True``) as columns.
+    s : (k,) array
+        Singular values.
+    VH : (k, n) array
+        Right singular vectors (if ``return_vecs=True``) as rows.
+    """
+    try:
+        return np.linalg.svd(A, full_matrices=False, compute_uv=return_vecs)
+
+    except np.linalg.linalg.LinAlgError:  # pragma: no cover
+        warnings.warn("Numpy SVD failed, trying again with different driver.")
+        return sla.svd(
+            A,
+            full_matrices=False,
+            compute_uv=return_vecs,
+            lapack_driver="gesvd",
+        )
+
+
+_SVDS_METHODS = {
+    "SLEPC": svds_slepc_spawn,
+    "SLEPC-NOMPI": svds_slepc,
+    "NUMPY": svds_numpy,
+    "SCIPY": svds_scipy,
+    "PRIMME": svds_primme,
+}
+
+
+def svds(A, k, ncv=None, return_vecs=True, backend="AUTO", **kwargs):
+    """Compute the partial singular value decomposition of an operator.
+
+    Parameters
+    ----------
+    A : dense, sparse or linear operator
+        The operator to decompose.
+    k : int, optional
+        number of singular value (triplets) to retrieve
+    ncv : int, optional
+        Number of lanczos vectors to use performing decomposition.
+    return_vecs : bool, optional
+        Whether to return the left and right vectors
+    backend : {'AUTO', 'SCIPY', 'SLEPC', 'SLEPC-NOMPI', 'NUMPY'}, optional
+        Which solver to use to perform decomposition.
+
+    Returns
+    -------
+    (Uk,) sk (, VHk) :
+        Singular value(s) (and vectors) such that ``Uk @ np.diag(sk) @ VHk``
+        approximates ``A``.
+    """
+    settings = {"k": k, "ncv": ncv, "return_vecs": return_vecs}
+
+    bkd = (
+        choose_backend(A, k, False)
+        if backend in {"auto", "AUTO"}
+        else backend.upper()
+    )
+    svds_func = _SVDS_METHODS[bkd.upper()]
+
+    return svds_func(A, **settings, **kwargs)
+
+
+# -------------------------------------------------------------------------- #
+# Norms and other quantities based on decompositions                         #
+# -------------------------------------------------------------------------- #
+
+
+def norm_2(A, **kwargs):
+    """Return the 2-norm of operator, ``A``, i.e. the largest singular value."""
+    return svds(A, k=1, return_vecs=False, **kwargs)[0]
+
+
+def norm_fro_dense(A):
+    """Frobenius norm for dense matrices"""
+    return vdot(A, A).real ** 0.5
+
+
+def norm_fro_sparse(A):
+    return vdot(A.data, A.data).real ** 0.5
+
+
+def norm_trace_dense(A, isherm=False):
+    """Returns the trace norm of operator ``A``, that is,
+    the sum of the absolute eigenvalues.
+    """
+    if isherm:
+        return abs(eigvalsh(A)).sum()
+    else:
+        return svd(A, return_vecs=False).sum()
+
+
+def norm(A, ntype=2, **kwargs):
+    """Operator norms.
+
+    Parameters
+    ----------
+    A : operator
+        The operator to find norm of.
+    ntype : str
+        Norm to calculate, if any of:
+
+        - {2, '2', 'spectral'}: largest singular value
+        - {'f', 'fro'}: frobenius norm
+        - {'t', 'nuc', 'tr', 'trace'}: sum of singular values
+
+    Returns
+    -------
+    x : float
+        The operator norm.
+    """
+    types = {
+        "2": "2",
+        2: "2",
+        "spectral": "2",
+        "f": "f",
+        "fro": "f",
+        "t": "t",
+        "trace": "t",
+        "nuc": "t",
+        "tr": "t",
+    }
+    methods = {
+        ("2", 0): norm_2,
+        ("2", 1): norm_2,
+        ("t", 0): norm_trace_dense,
+        ("f", 0): norm_fro_dense,
+        ("f", 1): norm_fro_sparse,
+    }
+    return methods[(types[ntype], issparse(A))](A, **kwargs)
+
+
+# --------------------------------------------------------------------------- #
+#                               Matrix functions                              #
+# --------------------------------------------------------------------------- #
+
+
+def expm(A, herm=False):
+    """Matrix exponential, can be accelerated if explicitly hermitian.
+
+    Parameters
+    ----------
+    A : dense or sparse operator
+        Operator to exponentiate.
+    herm : bool, optional
+        If True (not default), and ``A`` is dense, digonalize the matrix
+        in order to perform the exponential.
+    """
+    if issparse(A):
+        # convert to and from csc to suppress scipy warning
+        return spla.expm(A.tocsc()).tocsr()
+    elif not herm:
+        return qarray(spla.expm(A))
+    else:
+        evals, evecs = eigh(A)
+        return evecs @ ldmul(np.exp(evals), dag(evecs))
+
+
+_EXPM_MULTIPLY_METHODS = {
+    "SCIPY": spla.expm_multiply,
+    "SLEPC": functools.partial(mfn_multiply_slepc_spawn, fntype="exp"),
+    "SLEPC-KRYLOV": functools.partial(
+        mfn_multiply_slepc_spawn, fntype="exp", MFNType="KRYLOV"
+    ),
+    "SLEPC-EXPOKIT": functools.partial(
+        mfn_multiply_slepc_spawn, fntype="exp", MFNType="EXPOKIT"
+    ),
+    "SLEPC-NOMPI": functools.partial(mfn_multiply_slepc, fntype="exp"),
+}
+
+
+def expm_multiply(mat, vec, backend="AUTO", **kwargs):
+    """Compute the action of ``expm(mat)`` on ``vec``.
+
+    Parameters
+    ----------
+    mat : operator
+        Operator with which to act with exponential on ``vec``.
+    vec : vector-like
+        Vector to act with exponential of operator on.
+    backend : {'AUTO', 'SCIPY', 'SLEPC', 'SLEPC-KRYLOV', 'SLEPC-EXPOKIT'}
+        Which backend to use.
+    kwargs
+        Supplied to backend function.
+
+    Returns
+    -------
+    vector
+        Result of ``expm(mat) @ vec``.
+    """
+    if backend == "AUTO":
+        if SLEPC4PY_FOUND and vec.size > 2**10:
+            backend = "SLEPC"
+        else:
+            backend = "SCIPY"
+
+    return _EXPM_MULTIPLY_METHODS[backend.upper()](mat, vec, **kwargs)
+
+
+def sqrtm(A, herm=True):
+    """Matrix square root, can be accelerated if explicitly hermitian.
+
+    Parameters
+    ----------
+    A : dense array
+        Operator to take square root of.
+    herm : bool, optional
+        If True (the default), and ``A`` is dense, digonalize the matrix
+        in order to take the square root.
+
+    Returns
+    -------
+    array
+    """
+    if issparse(A):
+        raise NotImplementedError("No sparse sqrtm available.")
+    elif not herm:
+        return qarray(sla.sqrtm(A))
+    else:
+        evals, evecs = eigh(A)
+        return evecs @ ldmul(np.sqrt(evals.astype(complex)), dag(evecs))
+
+
+class IdentityLinearOperator(spla.LinearOperator):
+    """Get a ``LinearOperator`` representation of the identity operator,
+    scaled by ``factor``.
+
+    Parameters
+    ----------
+    size : int
+        The size of the identity.
+    factor : float
+        The coefficient of the identity.
+
+    Examples
+    --------
+
+    >>> I3 = IdentityLinearOperator(100, 1/3)
+    >>> p = rand_ket(100)
+    >>> np.allclose(I3 @ p, p / 3)
+    True
+    """
+
+    def __init__(self, size, factor=1):
+        self.factor = factor
+        super().__init__(dtype=np.array(factor).dtype, shape=(size, size))
+
+    def _matvec(self, vec):
+        return self.factor * vec
+
+    def _rmatvec(self, vec):
+        return self.factor * vec
+
+    def _matmat(self, mat):
+        return self.factor * mat
+
+
+class Lazy:
+    """A simple class representing an unconstructed matrix. This can be passed
+    to, for example, MPI workers, who can then construct the matrix themselves.
+    The main function ``fn`` should ideally take an ``ownership`` keyword to
+    avoid forming every row.
+
+    This is essentially like using ``functools.partial`` and assigning the
+    ``shape`` attribute.
+
+    Parameters
+    ----------
+    fn : callable
+        A function that constructs an operator.
+    shape :
+        Shape of the constructed operator.
+    args
+        Supplied to ``fn``.
+    kwargs
+        Supplied to ``fn``.
+
+    Returns
+    -------
+    Lazy : callable
+
+    Examples
+    --------
+    Setup the lazy operator:
+
+    >>> H_lazy = Lazy(ham_heis, n=10, shape=(2**10, 2**10), sparse=True)
+    >>> H_lazy
+    <Lazy(ham_heis, shape=(1024, 1024), dtype=None)>
+
+    Build a matrix slice (usually done automatically by e.g. ``eigs``):
+
+    >>> H_lazy(ownership=(256, 512))
+    <256x1024 sparse matrix of type '<class 'numpy.float64'>'
+            with 1664 stored elements in Compressed Sparse Row format>
+    """
+
+    def __init__(self, fn, *args, shape=None, factor=None, **kwargs):
+        if shape is None:
+            raise TypeError("`shape` must be specified.")
+        self.fn = fn
+        self.args = args
+        self.kwargs = kwargs
+        self.shape = shape
+        self.factor = factor
+        self.dtype = None
+
+    def __imul__(self, x):
+        if self.factor is None:
+            self.factor = x
+        else:
+            self.factor = self.factor * x
+
+    def __mul__(self, x):
+        if self.factor is not None:
+            x = x * self.factor
+        return Lazy(
+            self.fn, *self.args, shape=self.shape, factor=x, **self.kwargs
+        )
+
+    def __rmul__(self, x):
+        return self.__mul__(x)
+
+    def __call__(self, **kwargs):
+        A = self.fn(*self.args, **self.kwargs, **kwargs)
+
+        # check if any prefactors have been set
+        if self.factor is not None:
+            # try inplace first
+            try:
+                A *= self.factor
+            except (ValueError, TypeError):
+                A = self.factor * A
+
+        # helpful to store dtype once constructed
+        self.dtype = A.dtype
+        return A
+
+    def __repr__(self):
+        s = "<Lazy({}, shape={}{}{})>"
+
+        s_dtype = f", dtype={self.dtype}" if self.dtype is not None else ""
+        s_factor = f", factor={self.factor}" if self.factor is not None else ""
+
+        return s.format(self.fn.__name__, self.shape, s_dtype, s_factor)