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

trajectree/quimb/quimb/linalg/base_linalg.py

@@ -1,719 +0,0 @@
-"""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)