"""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 #
# --------------------------------------------------------------------------- #
[docs]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,
}
[docs]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 #
# --------------------------------------------------------------------------- #
[docs]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)
[docs]@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)
[docs]def groundstate(ham, **kwargs):
"""Alias for finding lowest eigenvector only.
"""
return eigvecsh(ham, k=1, which='SA', **kwargs)
[docs]def groundenergy(ham, **kwargs):
"""Alias for finding lowest eigenvalue only.
"""
return eigvalsh(ham, k=1, which='SA', **kwargs)[0]
[docs]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
[docs]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.A if issparse(A) else A, **kwargs)
else:
lk = eigvalsh(A.A 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]
[docs]def eigvalsh_window(*args, **kwargs):
"""Alias for only finding the eigenvalues in a relative window.
"""
return eigh_window(*args, return_vecs=False, **kwargs)
[docs]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 #
# -------------------------------------------------------------------------- #
[docs]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,
}
[docs]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 #
# -------------------------------------------------------------------------- #
[docs]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]
[docs]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
[docs]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()
[docs]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 #
# --------------------------------------------------------------------------- #
[docs]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'),
}
[docs]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)
[docs]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))
[docs]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
[docs]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)