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