"""Interface to slepc4py for solving advanced eigenvalue problems.
"""
import numpy as np
import scipy.sparse as sp
import quimb as qu
# --------------------------------------------------------------------------- #
# LAZY LOAD MPI/PETSc/SLEPc #
# --------------------------------------------------------------------------- #
[docs]def get_default_comm():
"""Define the default communicator.
"""
from mpi4py import MPI
return MPI.COMM_SELF
class _PetscSlepcHandler:
"""Helper class for caching the PETSc and SLEPc imports and making sure
that the MPI comm they use is aligned.
"""
def __init__(self):
self._comm = '__UNINITIALIZED__'
def init_petsc_slepc(self, comm):
import os
petsc_arch = os.environ.get("PETSC_ARCH", None)
import petsc4py
petsc4py.init(args=['-no_signal_handler'], arch=petsc_arch, comm=comm)
import slepc4py
slepc4py.init(args=['-no_signal_handler'], arch=petsc_arch)
self._comm = comm
self._PETSc = petsc4py.PETSc
self._SLEPc = slepc4py.SLEPc
def get_petsc(self, comm=None):
if comm is None:
comm = get_default_comm()
if comm != self._comm:
self.init_petsc_slepc(comm)
return self._PETSc, comm
def get_slepc(self, comm=None):
if comm is None:
comm = get_default_comm()
if comm != self._comm:
self.init_petsc_slepc(comm)
return self._SLEPc, comm
_petsc_slepc_handler = _PetscSlepcHandler()
def get_petsc(comm=None):
global _petsc_slepc_handler
return _petsc_slepc_handler.get_petsc(comm=comm)
def get_slepc(comm=None):
global _petsc_slepc_handler
return _petsc_slepc_handler.get_slepc(comm=comm)
# --------------------------------------------------------------------------- #
# PETSc FUNCTIONS #
# --------------------------------------------------------------------------- #
class PetscLinearOperatorContext:
def __init__(self, lo):
self.lo = lo
self.real = lo.dtype in (float, np.float_)
def mult(self, _, x, y):
y[:] = self.lo.matvec(x)
def multHermitian(self, _, x, y):
y[:] = self.lo.rmatvec(x)
def linear_operator_2_petsc_shell(lo, comm=None):
PETSc, comm = get_petsc(comm=comm)
context = PetscLinearOperatorContext(lo)
A = PETSc.Mat().createPython(lo.shape, context, comm=comm)
A.setUp()
return A
[docs]def slice_sparse_matrix_to_components(mat, ri, rf):
"""Slice the matrix `mat` between indices `ri` and `rf` -- for csr or bsr.
"""
return (mat.indptr[ri:rf + 1] - mat.indptr[ri],
mat.indices[mat.indptr[ri]:mat.indptr[rf]],
mat.data[mat.indptr[ri]:mat.indptr[rf]])
[docs]def convert_mat_to_petsc(mat, comm=None):
"""Convert a matrix to the relevant PETSc type, currently
only supports csr, bsr and dense matrices formats.
Parameters
----------
mat : dense, sparse, LinearOperator or Lazy matrix.
The operator to convert.
comm : mpi4py.MPI.Comm instance
The mpi communicator.
Returns
-------
pmat : petsc4py.PETSc.Mat
The matrix in petsc form - only the local part if running
across several mpi processes.
"""
if isinstance(mat, sp.linalg.LinearOperator):
return linear_operator_2_petsc_shell(mat, comm=comm)
PETSc, comm = get_petsc(comm=comm)
mpi_sz = comm.Get_size()
pmat = PETSc.Mat()
# retrieve before mat is possibly built into sliced matrix
shape = mat.shape
pmat.create(comm=comm)
pmat.setSizes(shape)
pmat.setFromOptions()
pmat.setUp()
ri, rf = pmat.getOwnershipRange()
# only consider the operator already sliced if owns whole
sliced = (mpi_sz == 1)
if isinstance(mat, qu.Lazy):
# operator hasn't been constructed yet
try:
# try and and lazily construct with slicing
mat = mat(ownership=(ri, rf))
sliced = True
except TypeError:
mat = mat()
# Sparse compressed or block row matrix
if sp.issparse(mat):
mat.sort_indices()
if sliced:
csr = (mat.indptr, mat.indices, mat.data)
else:
csr = slice_sparse_matrix_to_components(mat, ri, rf)
if sp.isspmatrix_csr(mat):
pmat.createAIJ(size=shape, nnz=mat.nnz, csr=csr, comm=comm)
elif sp.isspmatrix_bsr(mat):
pmat.createBAIJ(size=shape, bsize=mat.blocksize,
nnz=mat.nnz, csr=csr, comm=comm)
# Dense matrix
else:
if mpi_sz > 1 and not sliced:
pmat.createDense(size=shape, array=mat[ri:rf, :], comm=comm)
else:
pmat.createDense(size=shape, array=mat, comm=comm)
pmat.assemble()
return pmat
[docs]def convert_vec_to_petsc(vec, comm=None):
"""Convert a vector/array to the PETSc form.
Parameters
----------
vec : vector-like
Numpy array, will be unravelled to one dimension.
comm : mpi4py.MPI.Comm instance
The mpi communicator.
Returns
-------
pvec : petsc4py.PETSc.Vec
The vector in petsc form - only the local part if running
across several mpi processes.
"""
PETSc, comm = get_petsc(comm=comm)
mpi_sz = comm.Get_size()
pvec = PETSc.Vec()
# get shape before slicing
size = max(vec.shape)
pvec.create(comm=comm)
pvec.setSizes(size)
pvec.setFromOptions()
pvec.setUp()
ri, rf = pvec.getOwnershipRange()
# only consider the vector already sliced if owns whole
sliced = (mpi_sz == 1)
if isinstance(vec, qu.Lazy):
# vector hasn't been constructed yet
try:
# try and and lazily construct with slicing
vec = vec(ownership=(ri, rf))
sliced = True
except TypeError:
vec = vec()
array = np.asarray(vec).reshape(-1)
if not sliced:
array = array[ri:rf]
pvec.createWithArray(array, comm=comm)
return pvec
[docs]def new_petsc_vec(d, comm=None):
"""Create an empty petsc vector of size `d`.
Parameters
----------
d : int
Dimension of vector, i.e. the global size.
comm : mpi4py.MPI.Comm instance
The mpi communicator.
Returns
-------
pvec : petsc4py.PETSc.Vec
An empty vector in petsc form - only the local part if running
across several mpi processes.
"""
PETSc, comm = get_petsc(comm=comm)
pvec = PETSc.Vec()
pvec.create(comm=comm)
pvec.setSizes(d)
pvec.setFromOptions()
pvec.setUp()
return pvec
[docs]def gather_petsc_array(x, comm, out_shape=None):
"""Gather the petsc vector/matrix `x` to a single array on the master
process, assuming that owernership is sliced along the first dimension.
Parameters
----------
x : petsc4py.PETSc Mat or Vec
Distributed petsc array to gather.
comm : mpi4py.MPI.COMM
MPI communicator
out_shape : tuple, optional
If not None, reshape the output array to this.
Returns
-------
gathered : np.array master, None on workers (rank > 0)
"""
# get local numpy array
lx = x.getArray()
ox = np.empty(2, dtype=int)
ox[:] = x.getOwnershipRange()
# master only
if comm.Get_rank() == 0:
# create total array
ax = np.empty(x.getSize(), dtype=lx.dtype)
# set master's portion
ax[ox[0]:ox[1], ...] = lx
# get ownership ranges and data from worker processes
for i in range(1, comm.Get_size()):
comm.Recv(ox, source=i, tag=11)
# receive worker's part of ouput vector
comm.Recv(ax[ox[0]:ox[1], ...], source=i, tag=42)
if out_shape is not None:
ax = ax.reshape(*out_shape)
# Worker only
else:
# send ownership range
comm.Send(ox, dest=0, tag=11)
# send local portion of eigenvectors as buffer
comm.Send(lx, dest=0, tag=42)
ax = None
return ax
[docs]def normalize_real_part(vecs, transposed=False):
"""Take the real part of a set of vectors and normalize it. This is used
for returning real eigenvectors even when the PETSc scalar type is complex.
"""
k = vecs.shape[0] if transposed else vecs.shape[1]
vecs = np.ascontiguousarray(vecs.real)
for i in range(k):
where = (i, Ellipsis) if transposed else (Ellipsis, i)
qu.nmlz(vecs[where], inplace=True)
return vecs
# --------------------------------------------------------------------------- #
# SLEPc FUNCTIONS #
# --------------------------------------------------------------------------- #
def _init_krylov_subspace(comm=None, tol=None, maxiter=None,
KSPType="preonly",
PCType="lu",
PCFactorSolverType="mumps",
):
"""Initialise a krylov subspace and preconditioner.
"""
PETSc, comm = get_petsc(comm=comm)
K = PETSc.KSP().create(comm=comm)
K.setType(KSPType)
K.setTolerances(rtol=tol, max_it=maxiter)
if PCType:
PC = K.getPC()
PC.setType(PCType)
if PCFactorSolverType:
PC.setFactorSolverType(PCFactorSolverType)
PC.setFactorShift(PETSc.Mat.FactorShiftType.POSITIVE_DEFINITE)
PC.setFromOptions()
K.setPC(PC)
K.setFromOptions()
return K
def _init_spectral_inverter(STType="sinvert",
KSPType="preonly",
PCType="lu",
PCFactorSolverType="mumps",
comm=None):
"""Create a slepc spectral transformation object with specified solver.
"""
SLEPc, comm = get_slepc(comm=comm)
S = SLEPc.ST().create(comm=comm)
S.setType(STType)
# set the krylov subspace and preconditioner.
if KSPType:
K = _init_krylov_subspace(
KSPType=KSPType, PCType=PCType, comm=comm,
PCFactorSolverType=PCFactorSolverType)
S.setKSP(K)
S.setFromOptions()
return S
_WHICH_SCIPY_TO_SLEPC = {
"LM": 'LARGEST_MAGNITUDE',
"SM": 'SMALLEST_MAGNITUDE',
"LR": 'LARGEST_REAL',
"LA": 'LARGEST_REAL',
"SR": 'SMALLEST_REAL',
"SA": 'SMALLEST_REAL',
"LI": 'LARGEST_IMAGINARY',
"SI": 'SMALLEST_IMAGINARY',
"TM": 'TARGET_MAGNITUDE',
"TR": 'TARGET_REAL',
"TI": 'TARGET_IMAGINARY',
"ALL": 'ALL',
}
def _which_scipy_to_slepc(which):
SLEPc = get_slepc()[0]
return getattr(SLEPc.EPS.Which, _WHICH_SCIPY_TO_SLEPC[which.upper()])
def _init_eigensolver(k=6, which='LM', sigma=None, isherm=True, isgen=False,
EPSType=None, st_opts=None, tol=None, A_is_linop=False,
B_is_linop=False, maxiter=None, ncv=None, l_win=None,
comm=None):
"""Create an advanced eigensystem solver
Parameters
----------
sigma :
Target eigenvalue.
isherm :
Whether problem is hermitian or not.
Returns
-------
SLEPc solver ready to be called.
"""
SLEPc, comm = get_slepc(comm=comm)
eigensolver = SLEPc.EPS().create(comm=comm)
if st_opts is None:
st_opts = {}
if l_win is not None:
EPSType = {'ciss': 'ciss', None: 'ciss'}[EPSType]
which = 'ALL'
rg = eigensolver.getRG()
rg.setType(SLEPc.RG.Type.INTERVAL)
rg.setIntervalEndpoints(*l_win, -0.1, 0.1)
else:
eigensolver.setDimensions(k, ncv)
internal = (sigma is not None) or (l_win is not None)
# pick right backend
if EPSType is None:
if (isgen and B_is_linop) or (internal and A_is_linop):
EPSType = 'gd'
else:
EPSType = 'krylovschur'
# set some preconditioning defaults for 'gd' and 'lobpcg'
if EPSType in ('gd', 'lobpcg', 'blopex', 'primme'):
st_opts.setdefault('STType', 'precond')
st_opts.setdefault('KSPType', 'preonly')
st_opts.setdefault('PCType', 'none')
# set some preconditioning defaults for 'jd'
elif EPSType == 'jd':
st_opts.setdefault('STType', 'precond')
st_opts.setdefault('KSPType', 'bcgs')
st_opts.setdefault('PCType', 'none')
# set the spectral inverter / preconditioner.
if st_opts or internal:
st = _init_spectral_inverter(comm=comm, **st_opts)
eigensolver.setST(st)
# NB: `setTarget` must be called *after* `setST`.
if sigma is not None:
which = "TR"
eigensolver.setTarget(sigma)
if A_is_linop:
st.setMatMode(SLEPc.ST.MatMode.SHELL)
_EPS_PROB_TYPES = {
(False, False): SLEPc.EPS.ProblemType.NHEP,
(True, False): SLEPc.EPS.ProblemType.HEP,
(False, True): SLEPc.EPS.ProblemType.GNHEP,
(True, True): SLEPc.EPS.ProblemType.GHEP,
}
eigensolver.setType(EPSType)
eigensolver.setProblemType(_EPS_PROB_TYPES[(isherm, isgen)])
eigensolver.setWhichEigenpairs(_which_scipy_to_slepc(which))
eigensolver.setConvergenceTest(SLEPc.EPS.Conv.REL)
eigensolver.setTolerances(tol=tol, max_it=maxiter)
eigensolver.setFromOptions()
return eigensolver
[docs]def eigs_slepc(A, k, *, B=None, which=None, sigma=None, isherm=True, P=None,
v0=None, ncv=None, return_vecs=True, sort=True, EPSType=None,
return_all_conv=False, st_opts=None, tol=None, maxiter=None,
l_win=None, comm=None):
"""Solve a matrix using the advanced eigensystem solver
Parameters
----------
A : dense-matrix, sparse-matrix, LinearOperator or callable
Operator to solve.
k : int, optional
Number of requested eigenpairs.
B : dense-matrix, sparse-matrix, LinearOperator or callable
The RHS operator defining a generalized eigenproblem.
which : {"LM": "SM", "LR", "LA", "SR", "SA", "LI", "SI", "TM", "TR", "TI"}
Which eigenpairs to target. See :func:`scipy.sparse.linalg.eigs`.
sigma : float, optional
Target eigenvalue, implies ``which='TR'`` if this is not set but
``sigma`` is.
isherm : bool, optional
Whether problem is hermitian or not.
P : dense-matrix, sparse-matrix, LinearOperator or callable
Perform the eigensolve in the subspace defined by this projector.
v0 : 1D-array like, optional
Initial iteration vector, e.g., informed guess at eigenvector.
ncv : int, optional
Subspace size, defaults to ``min(20, 2 * k)``.
return_vecs : bool, optional
Whether to return the eigenvectors.
sort : bool, optional
Whether to sort the eigenpairs in ascending real value.
EPSType : {"krylovschur", 'gd', 'lobpcg', 'jd', 'ciss'}, optional
SLEPc eigensolver type to use, see slepc4py.EPSType.
return_all_conv : bool, optional
Whether to return converged eigenpairs beyond requested subspace size
st_opts : dict, optional
options to send to the eigensolver internal inverter.
tol : float, optional
Tolerance.
maxiter : int, optional
Maximum number of iterations.
comm : mpi4py communicator, optional
MPI communicator, defaults to ``COMM_SELF`` for a single process solve.
Returns
-------
lk : (k,) array
The eigenvalues.
vk : (m, k) array
Corresponding eigenvectors (if ``return_vecs=True``).
"""
if comm is None:
comm = get_default_comm()
A_is_linop = isinstance(A, sp.linalg.LinearOperator)
B_is_linop = isinstance(B, sp.linalg.LinearOperator)
isgen = B is not None
eigensolver = _init_eigensolver(
which=("SA" if (which is None) and (sigma is None) else
"TR" if (which is None) and (sigma is not None) else which),
EPSType=EPSType, k=k, sigma=sigma, isherm=isherm, tol=tol, ncv=ncv,
maxiter=maxiter, st_opts=st_opts, comm=comm, l_win=l_win, isgen=isgen,
A_is_linop=A_is_linop, B_is_linop=B_is_linop)
# set up the initial operators and solver
pA = convert_mat_to_petsc(A, comm=comm)
pB = convert_mat_to_petsc(B, comm=comm) if isgen else None
if P is not None:
pP = convert_mat_to_petsc(P, comm=comm)
pA = pP.transposeMatMult(pA.matMult(pP))
eigensolver.setOperators(pA, pB)
if v0 is not None:
eigensolver.setInitialSpace(convert_vec_to_petsc(v0, comm=comm))
eigensolver.solve()
# work out how many eigenpairs to retrieve
nconv = eigensolver.getConverged()
k = nconv if (return_all_conv or l_win is not None) else k
if nconv < k:
raise RuntimeError("SLEPC eigs did not find enough eigenpairs, "
f"wanted: {k}, found: {nconv}.")
# get eigenvalues
rank = comm.Get_rank()
if rank == 0:
lk = np.asarray([eigensolver.getEigenvalue(i) for i in range(k)])
lk = lk.real if isherm else lk
else:
res = None
# gather eigenvectors
if return_vecs:
pvec = pA.getVecLeft()
def get_vecs_local():
for i in range(k):
eigensolver.getEigenvector(i, pvec)
if P is not None:
pvecP = pP.getVecLeft()
pP.mult(pvec, pvecP)
yield gather_petsc_array(
pvecP if P is not None else pvec,
comm=comm, out_shape=(-1, 1))
lvecs = tuple(get_vecs_local())
if rank == 0:
vk = np.concatenate(lvecs, axis=1)
if sort:
sortinds = np.argsort(lk)
lk, vk = lk[sortinds], vk[:, sortinds]
# check if input matrix was real -> use real output when
# petsc compiled with complex scalars and thus outputs complex
convert = (isherm and np.issubdtype(A.dtype, np.floating) and
np.issubdtype(vk.dtype, np.complexfloating))
if convert:
vk = normalize_real_part(vk)
res = lk, qu.qarray(vk)
elif rank == 0:
res = np.sort(lk) if sort else lk
eigensolver.destroy()
return res
# ----------------------------------- SVD ----------------------------------- #
def _init_svd_solver(nsv=6, SVDType='cross', tol=None, maxiter=None,
ncv=None, comm=None):
SLEPc, comm = get_slepc(comm=comm)
svd_solver = SLEPc.SVD().create(comm=comm)
svd_solver.setType(SVDType)
svd_solver.setTolerances(tol=tol, max_it=maxiter)
svd_solver.setDimensions(nsv=nsv, ncv=ncv)
svd_solver.setFromOptions()
return svd_solver
[docs]def svds_slepc(A, k=6, ncv=None, return_vecs=True, SVDType='cross',
return_all_conv=False, tol=None, maxiter=None, comm=None):
"""Find the singular values for sparse matrix `a`.
Parameters
----------
A : sparse matrix in csr format
The matrix to solve.
k : int
Number of requested singular values.
method : {"cross", "cyclic", "lanczos", "trlanczos"}
Solver method to use.
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.
"""
if comm is None:
comm = get_default_comm()
pA = convert_mat_to_petsc(A, comm=comm)
svd_solver = _init_svd_solver(nsv=k, SVDType=SVDType, tol=tol,
maxiter=maxiter, ncv=ncv, comm=comm)
svd_solver.setOperator(pA)
svd_solver.solve()
nconv = svd_solver.getConverged()
k = nconv if return_all_conv else k
if nconv < k:
raise RuntimeError("SLEPC svds did not find enough singular triplets, "
f"wanted: {k}, found: {nconv}.")
rank = comm.Get_rank()
if return_vecs:
def usv_getter():
v, u = pA.createVecs()
for i in range(k):
s = svd_solver.getSingularTriplet(i, u, v)
lu = gather_petsc_array(u, comm=comm, out_shape=(-1, 1))
lv = gather_petsc_array(v, comm=comm, out_shape=(1, -1))
yield lu, s, lv
lus, sk, lvs = zip(*usv_getter())
sk = np.asarray(sk)
if rank == 0:
uk = qu.qarray(np.concatenate(lus, axis=1))
vtk = qu.qarray(np.concatenate(lvs, axis=0).conjugate())
# # check if input matrix was real -> use real output when
# # petsc compiled with complex scalars and thus outputs complex
# convert = (np.issubdtype(A.dtype, np.floating) and
# np.issubdtype(uk.dtype, np.complexfloating))
# if convert:
# uk = normalize_real_part(uk)
# vtk = normalize_real_part(uk, transposed=True)
res = uk, sk, vtk
else:
res = np.asarray([svd_solver.getValue(i) for i in range(k)])
svd_solver.destroy()
return res if rank == 0 else None
# ------------------------ matrix multiply function ------------------------- #
[docs]def mfn_multiply_slepc(mat, vec,
fntype='exp',
MFNType='AUTO',
comm=None,
isherm=False):
"""Compute the action of ``func(mat) @ vec``.
Parameters
----------
mat : operator
Operator to compute function action of.
vec : vector-like
Vector to compute matrix function action on.
func : {'exp', 'sqrt', 'log'}, optional
Function to use.
MFNType : {'krylov', 'expokit'}, optional
Method of computing the matrix function action, 'expokit' is only
available for func='exp'.
comm : mpi4py.MPI.Comm instance, optional
The mpi communicator.
isherm : bool, optional
If `mat` is known to be hermitian, this might speed things up in
some circumstances.
Returns
-------
fvec : array
The vector output of ``func(mat) @ vec``.
"""
SLEPc, comm = get_slepc(comm=comm)
mat = convert_mat_to_petsc(mat, comm=comm)
if isherm:
mat.setOption(mat.Option.HERMITIAN, True)
vec = convert_vec_to_petsc(vec, comm=comm)
out = new_petsc_vec(vec.size, comm=comm)
if MFNType.upper() == 'AUTO':
if (fntype == 'exp') and (vec.size <= 2**16):
MFNType = 'EXPOKIT'
else:
MFNType = 'KRYLOV'
# set up the matrix function options and objects
mfn = SLEPc.MFN().create(comm=comm)
mfn.setType(getattr(SLEPc.MFN.Type, MFNType.upper()))
mfn_fn = mfn.getFN()
mfn_fn.setType(getattr(SLEPc.FN.Type, fntype.upper()))
mfn_fn.setScale(1.0, 1.0)
mfn.setFromOptions()
mfn.setOperator(mat)
# 'solve' / perform the matrix function
mfn.solve(vec, out)
# --> gather the (distributed) petsc vector to a numpy matrix on master
all_out = gather_petsc_array(out, comm=comm, out_shape=(-1, 1))
mfn.destroy()
return all_out
# -------------------- solve linear system of equations --------------------- #
[docs]def lookup_ksp_error(i):
"""Look up PETSc error to print when raising after not converging.
"""
PETSc = get_petsc()
_KSP_DIVERGED_REASONS = {i: error for error, i in
PETSc.KSP.ConvergedReason.__dict__.items()
if isinstance(i, int)}
return _KSP_DIVERGED_REASONS[i]
def ssolve_slepc(A, y, isherm=True, comm=None, maxiter=None, tol=None,
KSPType='preonly',
PCType='lu',
PCFactorSolverType="mumps",
):
if comm is None:
comm = get_default_comm()
A = convert_mat_to_petsc(A, comm=comm)
if isherm:
A.setOption(A.Option.HERMITIAN, isherm)
x = A.createVecRight()
out_shape = y.shape
y = convert_vec_to_petsc(y, comm=comm)
ksp = _init_krylov_subspace(
KSPType=KSPType, PCType=PCType, comm=comm, maxiter=maxiter, tol=tol,
PCFactorSolverType=PCFactorSolverType)
ksp.setOperators(A)
ksp.solve(y, x)
converged_reason = ksp.getConvergedReason()
if converged_reason < 0:
raise RuntimeError("PETSc KSP solve did not converge, reason: "
f"{lookup_ksp_error(converged_reason)}")
x = gather_petsc_array(x, comm=comm, out_shape=out_shape)
ksp.destroy()
return x