def _lanczos_asis(a, V, u, alpha, beta, i_start, i_end):
for i in range(i_start, i_end):
u[...] = a @ V[i]
cublas.dotc(V[i], u, out=alpha[i])
u -= u.T @ V[:i + 1].conj().T @ V[:i + 1]
cublas.nrm2(u, out=beta[i])
if i >= i_end - 1:
break
V[i + 1] = u / beta[i]
def _update_asis(self, i_start, i_end):
for i in range(i_start, i_end):
u = self.A @ self.V[i]
cublas.dotc(self.V[i], u, out=self.alpha[i])
u -= u.T @ self.V[:i + 1].conj().T @ self.V[:i + 1]
cublas.nrm2(u, out=self.beta[i])
if i >= i_end - 1:
break
self.V[i + 1] = u / self.beta[i]
return u
def test_dotc(self):
x = self._make_random_vector()
y = self._make_random_vector()
ref = x.conj().dot(y)
out = self._make_out(self.dtype)
res = cublas.dotc(x, y, out=out)
self._check_pointer(res, out)
cupy.testing.assert_allclose(res, ref, rtol=self.tol, atol=self.tol)
def _lanczos_asis(a, V, u, alpha, beta, i_start, i_end):
beta_eps = inversion_eps(a.dtype)
for i in range(i_start, i_end):
u[...] = a @ V[i]
cublas.dotc(V[i], u, out=alpha[i])
u -= u.T @ V[:i + 1].conj().T @ V[:i + 1]
cublas.nrm2(u, out=beta[i])
if i >= i_end - 1:
break
if beta[i] < beta_eps:
V[i + 1:i_end, :] = 0
u[...] = 0
break
if i == i_start:
beta_eps *= beta[i] # scale eps to largest beta
V[i + 1] = u / beta[i]
def eigsh(a,
k=6,
*,
which='LM',
ncv=None,
maxiter=None,
tol=0,
return_eigenvectors=True):
"""Finds ``k`` eigenvalues and eigenvectors of the real symmetric matrix.
Solves ``Ax = wx``, the standard eigenvalue problem for ``w`` eigenvalues
with corresponding eigenvectors ``x``.
Args:
a (ndarray, spmatrix or LinearOperator): A symmetric square matrix with
dimension ``(n, n)``. ``a`` must :class:`cupy.ndarray`,
:class:`cupyx.scipy.sparse.spmatrix` or
:class:`cupyx.scipy.sparse.linalg.LinearOperator`.
k (int): The number of eigenvalues and eigenvectors to compute. Must be
``1 <= k < n``.
which (str): 'LM' or 'LA'. 'LM': finds ``k`` largest (in magnitude)
eigenvalues. 'LA': finds ``k`` largest (algebraic) eigenvalues.
ncv (int): The number of Lanczos vectors generated. Must be
``k + 1 < ncv < n``. If ``None``, default value is used.
maxiter (int): Maximum number of Lanczos update iterations.
If ``None``, default value is used.
tol (float): Tolerance for residuals ``||Ax - wx||``. If ``0``, machine
precision is used.
return_eigenvectors (bool): If ``True``, returns eigenvectors in
addition to eigenvalues.
Returns:
tuple:
If ``return_eigenvectors is True``, it returns ``w`` and ``x``
where ``w`` is eigenvalues and ``x`` is eigenvectors. Otherwise,
it returns only ``w``.
.. seealso:: :func:`scipy.sparse.linalg.eigsh`
.. note::
This function uses the thick-restart Lanczos methods
(https://sdm.lbl.gov/~kewu/ps/trlan.html).
"""
n = a.shape[0]
if a.ndim != 2 or a.shape[0] != a.shape[1]:
raise ValueError('expected square matrix (shape: {})'.format(a.shape))
if a.dtype.char not in 'fdFD':
raise TypeError('unsupprted dtype (actual: {})'.format(a.dtype))
if k <= 0:
raise ValueError('k must be greater than 0 (actual: {})'.format(k))
if k >= n:
raise ValueError('k must be smaller than n (actual: {})'.format(k))
if which not in ('LM', 'LA'):
raise ValueError('which must be \'LM\' or \'LA\' (actual: {})'
''.format(which))
if ncv is None:
ncv = min(max(2 * k, k + 32), n - 1)
else:
ncv = min(max(ncv, k + 2), n - 1)
if maxiter is None:
maxiter = 10 * n
if tol == 0:
tol = numpy.finfo(a.dtype).eps
alpha = cupy.zeros((ncv, ), dtype=a.dtype)
beta = cupy.zeros((ncv, ), dtype=a.dtype.char.lower())
V = cupy.empty((ncv, n), dtype=a.dtype)
# Set initial vector
u = cupy.random.random((n, )).astype(a.dtype)
V[0] = u / cublas.nrm2(u)
# Choose Lanczos implementation, unconditionally use 'fast' for now
upadte_impl = 'fast'
if upadte_impl == 'fast':
lanczos = _lanczos_fast(a, n, ncv)
else:
lanczos = _lanczos_asis
# Lanczos iteration
lanczos(a, V, u, alpha, beta, 0, ncv)
iter = ncv
w, s = _eigsh_solve_ritz(alpha, beta, None, k, which)
x = V.T @ s
# Compute residual
beta_k = beta[-1] * s[-1, :]
res = cublas.nrm2(beta_k)
while res > tol and iter < maxiter:
# Setup for thick-restart
beta[:k] = 0
alpha[:k] = w
V[:k] = x.T
u -= u.T @ V[:k].conj().T @ V[:k]
V[k] = u / cublas.nrm2(u)
u[...] = a @ V[k]
cublas.dotc(V[k], u, out=alpha[k])
u -= alpha[k] * V[k]
u -= V[:k].T @ beta_k
cublas.nrm2(u, out=beta[k])
V[k + 1] = u / beta[k]
# Lanczos iteration
lanczos(a, V, u, alpha, beta, k + 1, ncv)
iter += ncv - k
w, s = _eigsh_solve_ritz(alpha, beta, beta_k, k, which)
x = V.T @ s
# Compute residual
beta_k = beta[-1] * s[-1, :]
res = cublas.nrm2(beta_k)
if return_eigenvectors:
idx = cupy.argsort(w)
return w[idx], x[:, idx]
else:
return cupy.sort(w)
def cg(A,
b,
x0=None,
tol=1e-5,
maxiter=None,
M=None,
callback=None,
atol=None):
"""Uses Conjugate Gradient iteration to solve ``Ax = b``.
Args:
A (ndarray, spmatrix or LinearOperator): The real or complex matrix of
the linear system with shape ``(n, n)``. ``A`` must be a hermitian,
positive definitive matrix with type of :class:`cupy.ndarray`,
:class:`cupyx.scipy.sparse.spmatrix` or
:class:`cupyx.scipy.sparse.linalg.LinearOperator`.
b (cupy.ndarray): Right hand side of the linear system with shape
``(n,)`` or ``(n, 1)``.
x0 (cupy.ndarray): Starting guess for the solution.
tol (float): Tolerance for convergence.
maxiter (int): Maximum number of iterations.
M (ndarray, spmatrix or LinearOperator): Preconditioner for ``A``.
The preconditioner should approximate the inverse of ``A``.
``M`` must be :class:`cupy.ndarray`,
:class:`cupyx.scipy.sparse.spmatrix` or
:class:`cupyx.scipy.sparse.linalg.LinearOperator`.
callback (function): User-specified function to call after each
iteration. It is called as ``callback(xk)``, where ``xk`` is the
current solution vector.
atol (float): Tolerance for convergence.
Returns:
tuple:
It returns ``x`` (cupy.ndarray) and ``info`` (int) where ``x`` is
the converged solution and ``info`` provides convergence
information.
.. seealso:: :func:`scipy.sparse.linalg.cg`
"""
A, M, x, b = _make_system(A, M, x0, b)
matvec = A.matvec
psolve = M.matvec
n = A.shape[0]
if maxiter is None:
maxiter = n * 10
if n == 0:
return cupy.empty_like(b), 0
b_norm = cupy.linalg.norm(b)
if b_norm == 0:
return b, 0
if atol is None:
atol = tol * float(b_norm)
else:
atol = max(float(atol), tol * float(b_norm))
r = b - matvec(x)
iters = 0
rho = 0
while iters < maxiter:
z = psolve(r)
rho1 = rho
rho = cublas.dotc(r, z)
if iters == 0:
p = z
else:
beta = rho / rho1
p = z + beta * p
q = matvec(p)
alpha = rho / cublas.dotc(p, q)
x = x + alpha * p
r = r - alpha * q
iters += 1
if callback is not None:
callback(x)
resid = cublas.nrm2(r)
if resid <= atol:
break
info = 0
if iters == maxiter and not (resid <= atol):
info = iters
return x, info
def cg(A,
b,
x0=None,
tol=1e-5,
maxiter=None,
M=None,
callback=None,
atol=None):
"""Uses Conjugate Gradient iteration to solve ``Ax = b``.
Args:
A (cupy.ndarray or cupyx.scipy.sparse.spmatrix): The real or complex
matrix of the linear system with shape ``(n, n)``. ``A`` must
be a hermitian, positive definitive matrix.
b (cupy.ndarray): Right hand side of the linear system with shape
``(n,)`` or ``(n, 1)``.
x0 (cupy.ndarray): Starting guess for the solution.
tol (float): Tolerance for convergence.
maxiter (int): Maximum number of iterations.
M (cupy.ndarray or cupyx.scipy.sparse.spmatrix): Preconditioner for
``A``. The preconditioner should approximate the inverse of ``A``.
callback (function): User-specified function to call after each
iteration. It is called as ``callback(xk)``, where ``xk`` is the
current solution vector.
atol (float): Tolerance for convergence.
Returns:
tuple:
It returns ``x`` (cupy.ndarray) and ``info`` (int) where ``x`` is
the converged solution and ``info`` provides convergence
information.
.. seealso:: :func:`scipy.sparse.linalg.cg`
"""
if A.ndim != 2 or A.shape[0] != A.shape[1]:
raise ValueError('expected square matrix (shape: {})'.format(A.shape))
if A.dtype.char not in 'fdFD':
raise TypeError('unsupprted dtype (actual: {})'.format(A.dtype))
n = A.shape[0]
if not (b.shape == (n, ) or b.shape == (n, 1)):
raise ValueError('b has incompatible dimensins')
b = b.astype(A.dtype).ravel()
if n == 0:
return cupy.empty_like(b), 0
b_norm = cupy.linalg.norm(b)
if b_norm == 0:
return b, 0
if atol is None:
atol = tol * float(b_norm)
else:
atol = max(float(atol), tol * float(b_norm))
if x0 is None:
x = cupy.zeros((n, ), dtype=A.dtype)
else:
if not (x0.shape == (n, ) or x0.shape == (n, 1)):
raise ValueError('x0 has incompatible dimensins')
x = x0.astype(A.dtype).ravel()
if maxiter is None:
maxiter = n * 10
matvec, psolve = _make_funcs(A, M)
r = b - matvec(x)
iters = 0
rho = 0
while iters < maxiter:
z = psolve(r)
rho1 = rho
rho = cublas.dotc(r, z)
if iters == 0:
p = z
else:
beta = rho / rho1
p = z + beta * p
q = matvec(p)
alpha = rho / cublas.dotc(p, q)
x = x + alpha * p
r = r - alpha * q
iters += 1
if callback is not None:
callback(x)
resid = cublas.nrm2(r)
if resid <= atol:
break
info = 0
if iters == maxiter and not (resid <= atol):
info = iters
return x, info
