diff --git a/docs/source/gpu.rst b/docs/source/gpu.rst index 1229b0a2..5b5c75b0 100755 --- a/docs/source/gpu.rst +++ b/docs/source/gpu.rst @@ -226,7 +226,7 @@ GPU with CuPy, and GPU with JAX): - GPU/TPU with JAX * - :meth:`pylops.LinearOperator.cond` - |:white_check_mark:| - - |:red_circle:| + - |:white_check_mark:| - |:red_circle:| * - :meth:`pylops.LinearOperator.conj` - |:white_check_mark:| @@ -238,7 +238,7 @@ GPU with CuPy, and GPU with JAX): - |:white_check_mark:| * - :meth:`pylops.LinearOperator.eigs` - |:white_check_mark:| - - |:red_circle:| + - |:white_check_mark:| - |:red_circle:| * - :meth:`pylops.LinearOperator.todense` - |:white_check_mark:| @@ -632,12 +632,19 @@ Medical: .. warning:: - 1. The JAX backend of the :class:`pylops.signalprocessing.Convolve1D` operator + 1. The CuPy backend of :meth:`pylops.LinearOperator.eigs` does not support + non-symmetric operators, since CuPy does not provide an equivalent + of :func:`scipy.sparse.linalg.eigs`. A :obj:`NotImplementedError` is raised. + + 2. The CuPy backend of :meth:`pylops.LinearOperator.cond` can only be computed + with the ``uselobpcg=True`` option (and therefore for symmetric operators). + + 3. The JAX backend of the :class:`pylops.signalprocessing.Convolve1D` operator currently works only with 1d-arrays due to a different behaviour of :meth:`scipy.signal.convolve` and :meth:`jax.scipy.signal.convolve` with nd-arrays. - 2. The JAX backend of the :class:`pylops.avo.prestack.PrestackLinearModelling` + 4. The JAX backend of the :class:`pylops.avo.prestack.PrestackLinearModelling` operator currently works only with ``explicit=True`` due to the same issue as in point 1 for the :class:`pylops.signalprocessing.Convolve1D` operator employed when ``explicit=False``. diff --git a/pylops/linearoperator.py b/pylops/linearoperator.py index 3157cf3b..f1037da7 100644 --- a/pylops/linearoperator.py +++ b/pylops/linearoperator.py @@ -5,6 +5,7 @@ "aslinearoperator", ] +import warnings from abc import ABC, abstractmethod from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor @@ -14,13 +15,11 @@ import numpy as np import scipy as sp from numpy.linalg import solve as np_solve -from scipy.linalg import eigvals, lstsq +from scipy.linalg import lstsq from scipy.linalg import solve as sp_solve from scipy.sparse import csr_matrix from scipy.sparse.linalg import LinearOperator as spLinearOperator from scipy.sparse.linalg import eigs as sp_eigs -from scipy.sparse.linalg import eigsh as sp_eigsh -from scipy.sparse.linalg import lobpcg as sp_lobpcg from scipy.sparse.linalg import lsqr, spsolve # need to check scipy version since the interface submodule changed into @@ -35,10 +34,24 @@ from pylops._multioperator import _matvec_rmatvec_map from pylops.optimization.basic import cgls from pylops.utils._internal import _get_dtype -from pylops.utils.backend import get_array_module, get_module, get_sparse_eye +from pylops.utils.backend import ( + get_array_module, + get_eigsh, + get_eigvals, + get_lobpcg, + get_module, + get_sparse_eye, +) from pylops.utils.decorators import count from pylops.utils.estimators import trace_hutchinson, trace_hutchpp, trace_nahutchpp -from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray, ShapeLike, Tpool +from pylops.utils.typing import ( + DTypeLike, + InputDimsLike, + NDArray, + ShapeLike, + Tbackend, + Tpool, +) class _LinearOperator(ABC): @@ -178,6 +191,9 @@ def __init__( self.forceflat = forceflat self.name = name + # ndim: required by some eigenvalue solvers + self.ndim = len(self.shape) + # counters self.matvec_count = 0 self.rmatvec_count = 0 @@ -1065,31 +1081,47 @@ def eigs( symmetric: bool = False, niter: int | None = None, uselobpcg: bool = False, + backend: Tbackend = "numpy", **kwargs_eig: int | float | str, ) -> NDArray: - r"""Most significant eigenvalues of linear operator. + r"""N-most significant eigenvalues of a linear operator. - Return an estimate of the most significant eigenvalues + Return an estimate of the ``neigs`` most significant eigenvalues of the linear operator. If the operator has rectangular - shape (``shape[0]!=shape[1]``), eigenvalues are first + shape (``shape[0] != shape[1]``), eigenvalues are first computed for the square operator :math:`\mathbf{A^H}\mathbf{A}` and the square-root values are returned. + .. note:: + If ``backend="cupy"``, the eigenvalues of a square operator + can be computed only when in symmetric (``symmetric=True``), + since CuPy does not provide an equivalent + of :func:`scipy.sparse.linalg.eigs`. A + :obj:`NotImplementedError` is raised. + Parameters ---------- neigs : :obj:`int` - Number of eigenvalues to compute (if ``None``, return all). Note - that for ``explicit=False``, only :math:`N-1` eigenvalues can be - computed where :math:`N` is the size of the operator in the - model space + Number of eigenvalues to compute. Note that for ``explicit=False``, + only :math:`N-1` eigenvalues can be computed for real-valued + symmetric or non-square operators, and :math:`N-2` eigenvalues + otherwise, where :math:`N` is the size of the operator in the + model space. If ``None``, return :math:`N` for explicit case and + :math:`N-2` for implicit operators. symmetric : :obj:`bool`, optional - Operator is symmetric (``True``) or not (``False``). User should - set this parameter to ``True`` only when it is guaranteed that the - operator is real-symmetric or complex-hermitian matrices + Operator is real-symmetric or complex-hermitian (``True``) + or not (``False``). User should set this parameter to ``True`` + only when it is guaranteed that the operator is real-symmetric + or complex-hermitian niter : :obj:`int`, optional Number of iterations for eigenvalue estimation uselobpcg : :obj:`bool`, optional - Use :func:`scipy.sparse.linalg.lobpcg` + Use :func:`scipy.sparse.linalg.lobpcg`. Note that this method + is only available for symmetric operators. + backend : :obj:`str`, optional + .. versionadded:: 2.9.0 + + Backend to use (`numpy` or `cupy`) **kwargs_eig Arbitrary keyword arguments for :func:`scipy.sparse.linalg.eigs`, :func:`scipy.sparse.linalg.eigsh`, or @@ -1104,7 +1136,16 @@ def eigs( ------ ValueError If ``uselobpcg=True`` for a non-symmetric square matrix with - complex type + complex type, or if ``uselobpcg=True`` and the size of the operator + minus the number of constraints ``Y`` is smaller than 5 times ``neigs`` + (in which case :func:`scipy.sparse.linalg.lobpcg` would switch to a + dense solver) + NotImplementedError + If ``backend="cupy"`` and ``symmetric=False`` for a square + operator for which only a limited number of eigenvalues is + requested (``neigs`` smaller than the size of the operator, or + ``explicit=False``), as CuPy does not provide an equivalent of + :func:`scipy.sparse.linalg.eigs` Notes ----- @@ -1117,8 +1158,12 @@ def eigs( ``symmetric=True`` and the :func:`scipy.sparse.linalg.eigs` method is used ``symmetric=False``. However, when the matrix is represented explicitly within the linear operator (``explicit=True``) and all the - eigenvalues are requested the :func:`scipy.linalg.eigvals` is used - instead. + eigenvalues are requested (or ``neigs`` is equal to the number of + columns of the operator minus one for a square, non-symmetric operator) the + :func:`scipy.linalg.eigvals` is used instead. Note that for + or a complex-valued operator, when ``neigs`` is equal to the number of + columns of the operator minus one, the eigenvalues are selected based on + ``which`` and ``sigma`` is ignored. Finally, when only a limited number of eigenvalues is required, it is also possible to explicitly choose to use the @@ -1134,88 +1179,122 @@ def eigs( .. [1] http://www.caam.rice.edu/software/ARPACK/ """ - if self.explicit and isinstance(self.A, np.ndarray): - if self.shape[0] == self.shape[1]: - if neigs is None or neigs == self.shape[1]: - eigenvalues = eigvals(self.A) - else: - if not symmetric and np.iscomplexobj(self) and uselobpcg: - msg = ( - "Cannot use scipy.sparse.linalg.lobpcg " - "for non-symmetric square matrices of complex type..." - ) - raise ValueError(msg) - if symmetric and uselobpcg: - X = np.random.rand(self.shape[0], neigs).astype(self.dtype) - eigenvalues = sp_lobpcg( - self.A, X=X, maxiter=niter, **kwargs_eig - )[0] - elif symmetric: - eigenvalues = sp_eigsh( - self.A, k=neigs, maxiter=niter, **kwargs_eig - )[0] - else: - eigenvalues = sp_eigs( - self.A, k=neigs, maxiter=niter, **kwargs_eig - )[0] + def _check_lobpcg_size(n: int, neigs: int, Y: NDArray | None = None) -> None: + """Check that scipy's lobpcg will iterate rather than use a dense solver + + :func:`scipy.sparse.linalg.lobpcg` falls back to a dense eigensolver + (which cannot be applied to PyLops operators) when + ``n - sizeY < 5 * neigs``, where ``sizeY`` is the number of + constraints (columns of ``Y``, which can be passed via ``kwargs_eig``). + + """ + sizeY = Y.shape[1] if Y is not None and Y.ndim == 2 else 0 + if n - sizeY < 5 * neigs: + msg = ( + f"uselobpcg=True requires the operator size ({n}) minus the " + f"number of constraints ({sizeY}) to be at least 5 times the " + f"number of requested eigenvalues ({neigs}); reduce neigs to " + f"at most {(n - sizeY) // 5} or use uselobpcg=False" + ) + raise ValueError(msg) + + # Check if lobpcgc is feasible, raise error otherwise + square = self.shape[0] == self.shape[1] + if square and not symmetric and uselobpcg: + msg = ( + "Cannot use scipy.sparse.linalg.lobpcg " + "for non real-symmetric or complex-hermitian operators..." + ) + raise ValueError(msg) + + ncp = get_module(backend) + explicit = self.explicit and isinstance(self.A, ncp.ndarray) + if neigs is None: + neigs = self.shape[1] if explicit else self.shape[1] - 2 + # scipy.sparse.linalg.eigs (also used by eigsh for complex operators) + # requires neigs < N - 1, use dense solver instead + useeigs = (square and not symmetric) or np.iscomplexobj( + np.ones(1, dtype=self.dtype) + ) + # Check if all eigenvalues should be computed + alleigs = explicit and ( + neigs >= self.shape[1] or (useeigs and neigs == self.shape[1] - 1) + ) + + # Check that the number of eigenvalues requested is not larger + # than the number of columns in the operator minus two (eigs) or + # minus one (eigsh) + neigsmax = self.shape[1] - 2 if useeigs else self.shape[1] - 1 + if not explicit and neigs > neigsmax: + warnings.warn( + f"The provided neigs ({neigs}) exceeds the number of columns " + f"in the operator minus {self.shape[1] - neigsmax}, " + f"reset to {neigsmax}", + stacklevel=2, + ) + neigs = neigsmax + + if square: + # Square operator + Op = self.A if explicit else self + if alleigs: + # All eigenvalues (only reached for explicit=True) + eigenvalues = get_eigvals(ncp.ones(1))(Op) else: - if neigs is None or neigs == self.shape[1]: - eigenvalues = np.sqrt(eigvals(np.dot(np.conj(self.A.T), self.A))) - else: - if uselobpcg: - X = np.random.rand(self.shape[1], neigs).astype(self.dtype) - eigenvalues = np.sqrt( - sp_lobpcg( - np.dot(np.conj(self.A.T), self.A), - X=X, - maxiter=niter, - **kwargs_eig, - )[0] - ) - else: - eigenvalues = np.sqrt( - sp_eigsh( - np.dot(np.conj(self.A.T), self.A), - k=neigs, - maxiter=niter, - **kwargs_eig, - )[0] - ) - else: - if neigs is None or neigs >= self.shape[1]: - neigs = self.shape[1] - 2 - if self.shape[0] == self.shape[1]: - if not symmetric and np.iscomplexobj(self) and uselobpcg: - msg = ( - "Cannot use scipy.sparse.linalg.lobpcg for " - "non symmetric square matrices of complex type..." - ) - raise ValueError(msg) + # Limited number of eigenvalues if symmetric and uselobpcg: - X = np.random.rand(self.shape[0], neigs).astype(self.dtype) - eigenvalues = sp_lobpcg(self, X=X, maxiter=niter, **kwargs_eig)[0] + if not explicit: + _check_lobpcg_size(self.shape[0], neigs, kwargs_eig.get("Y")) + X = ncp.random.rand(self.shape[0], neigs).astype(self.dtype) + eigenvalues = get_lobpcg(ncp.ones(1))( + Op, X=X, maxiter=niter, **kwargs_eig + )[0] elif symmetric: - eigenvalues = sp_eigsh(self, k=neigs, maxiter=niter, **kwargs_eig)[ - 0 - ] + eigenvalues = get_eigsh(ncp.ones(1))( + Op, k=neigs, maxiter=niter, **kwargs_eig + )[0] else: - eigenvalues = sp_eigs(self, k=neigs, maxiter=niter, **kwargs_eig)[0] + if backend == "cupy": + msg = ( + "Eigenvalues of non-symmetric square " + "operators are not available with CuPy backend" + ) + raise NotImplementedError(msg) + eigenvalues = sp_eigs(Op, k=neigs, maxiter=niter, **kwargs_eig)[0] + else: + # Non-square operator (square-root of eigenvalues of A^H A) + Op = ncp.dot(ncp.conj(self.A.T), self.A) if explicit else self.H * self + if alleigs: + # All eigenvalues (only reached for explicit=True) + eigenvalues = get_eigvals(ncp.ones(1))(Op) + elif uselobpcg: + # Limited number of eigenvalues + if not explicit: + _check_lobpcg_size(self.shape[1], neigs, kwargs_eig.get("Y")) + X = ncp.random.rand(self.shape[1], neigs).astype(self.dtype) + eigenvalues = get_lobpcg(ncp.ones(1))( + Op, X=X, maxiter=niter, **kwargs_eig + )[0] else: - if uselobpcg: - X = np.random.rand(self.shape[1], neigs).astype(self.dtype) - eigenvalues = np.sqrt( - sp_lobpcg(self.H * self, X=X, maxiter=niter, **kwargs_eig)[0] - ) - else: - eigenvalues = np.sqrt( - sp_eigs(self.H * self, k=neigs, maxiter=niter, **kwargs_eig)[0] - ) - return -np.sort(-eigenvalues) + # Limited number of eigenvalues + eigenvalues = get_eigsh(ncp.ones(1))( + Op, k=neigs, maxiter=niter, **kwargs_eig + )[0] + eigenvalues = ncp.sqrt(eigenvalues) + if alleigs and neigs == self.shape[1] - 1: + # dense solver returns all eigenvalues: select neigs based on which + which = kwargs_eig.get("which", "LM") + key = {"M": ncp.abs, "R": ncp.real, "A": ncp.real, "I": ncp.imag}[which[1]] + key = key(eigenvalues) + idx = ncp.argsort(key if which[0] == "S" else -key)[:neigs] + eigenvalues = eigenvalues[idx] + return -ncp.sort(-eigenvalues) def cond( self, uselobpcg: bool = False, + backend: Tbackend = "numpy", **kwargs_eig: int | float | str, ) -> NDArray: r"""Condition number of linear operator. @@ -1223,10 +1302,19 @@ def cond( Return an estimate of the condition number of the linear operator as the ratio of the largest and lowest estimated eigenvalues. + .. note:: + If ``backend="cupy"``, the condition number can only be computed + with the ``uselobpcg=True`` option (and therefore for square, + symmetric operators only). + Parameters ---------- uselobpcg : :obj:`bool`, optional Use :func:`scipy.sparse.linalg.lobpcg` to compute eigenvalues + backend : :obj:`str`, optional + .. versionadded:: 2.9.0 + + Backend to use (`numpy` or `cupy`). **kwargs_eig Arbitrary keyword arguments for :func:`scipy.sparse.linalg.eigs`, :func:`scipy.sparse.linalg.eigsh`, or @@ -1258,15 +1346,27 @@ def cond( """ if not uselobpcg: - cond = ( - self.eigs(neigs=1, which="LM", **kwargs_eig).item() - / self.eigs(neigs=1, which="SM", **kwargs_eig).item() - ) + lmax = self.eigs(neigs=1, which="LM", backend=backend, **kwargs_eig).item() + lmin = self.eigs(neigs=1, which="SM", backend=backend, **kwargs_eig).item() + cond = lmax / lmin else: - cond = ( - self.eigs(neigs=1, uselobpcg=True, largest=True, **kwargs_eig).item() - / self.eigs(neigs=1, uselobpcg=True, largest=False, **kwargs_eig).item() - ) + lmax = self.eigs( + neigs=1, + symmetric=True, + uselobpcg=True, + largest=True, + backend=backend, + **kwargs_eig, + ).item() + lmin = self.eigs( + neigs=1, + symmetric=True, + uselobpcg=True, + largest=False, + backend=backend, + **kwargs_eig, + ).item() + cond = lmax / lmin return cond diff --git a/pylops/optimization/eigs.py b/pylops/optimization/eigs.py index 13e6412f..fd187d4c 100644 --- a/pylops/optimization/eigs.py +++ b/pylops/optimization/eigs.py @@ -1,15 +1,18 @@ __all__ = ["power_iteration"] +from typing import TYPE_CHECKING import numpy as np -from pylops import LinearOperator from pylops.utils.backend import get_module from pylops.utils.typing import NDArray, Tbackend +if TYPE_CHECKING: + from pylops.linearoperator import LinearOperator + def power_iteration( - Op: LinearOperator, + Op: "LinearOperator", niter: int = 10, tol: float = 1e-5, dtype: str = "float32", diff --git a/pylops/signalprocessing/_baseffts.py b/pylops/signalprocessing/_baseffts.py index 2806f6db..c151eec7 100644 --- a/pylops/signalprocessing/_baseffts.py +++ b/pylops/signalprocessing/_baseffts.py @@ -42,11 +42,11 @@ def __init__( dims = _value_or_sized_to_array(dims) _raise_on_wrong_dtype(dims, np.integer, "dims") - self.ndim = len(dims) + self.ndims = len(dims) axes = _value_or_sized_to_array(axis) _raise_on_wrong_dtype(axes, np.integer, "axis") - self.axis = get_normalize_axis_index()(axes[0], self.ndim) + self.axis = get_normalize_axis_index()(axes[0], self.ndims) if nfft is None: self.nfft = dims[self.axis] @@ -63,7 +63,7 @@ def __init__( self.doifftpad = False if self.nfft < dims[self.axis]: self.doifftpad = True - self.ifftpad = [(0, 0)] * self.ndim + self.ifftpad = [(0, 0)] * self.ndims self.ifftpad[self.axis] = (0, dims[self.axis] - self.nfft) warnings.warn( f"nfft={self.nfft} has been selected to be smaller than the size of " @@ -148,11 +148,11 @@ def __init__( dims = _value_or_sized_to_array(dims) _raise_on_wrong_dtype(dims, np.integer, "dims") - self.ndim = len(dims) + self.ndims = len(dims) axes = _value_or_sized_to_array(axes) _raise_on_wrong_dtype(axes, np.integer, "axes") - self.axes = np.array([get_normalize_axis_index()(d, self.ndim) for d in axes]) + self.axes = np.array([get_normalize_axis_index()(d, self.ndims) for d in axes]) self.naxes = len(self.axes) if self.naxes != len(np.unique(self.axes)): warnings.warn( @@ -210,7 +210,7 @@ def __init__( ] self.doifftpad = any(nfftshort) if self.doifftpad: - self.ifftpad = [(0, 0)] * self.ndim + self.ifftpad = [(0, 0)] * self.ndims for idir, (direction, nfshort) in enumerate( zip(self.axes, nfftshort, strict=True) ): diff --git a/pylops/signalprocessing/dtcwt.py b/pylops/signalprocessing/dtcwt.py index 6e9a24fa..66d803b8 100644 --- a/pylops/signalprocessing/dtcwt.py +++ b/pylops/signalprocessing/dtcwt.py @@ -109,10 +109,9 @@ def __init__( if dtcwt_message is not None: raise NotImplementedError(dtcwt_message) - dims = _value_or_sized_to_tuple(dims) - self.ndim = len(dims) self.axis = axis + dims = _value_or_sized_to_tuple(dims) self.otherdims = int(np.prod(dims) / dims[self.axis]) self.dims_swapped = list(dims) self.dims_swapped[0], self.dims_swapped[self.axis] = ( diff --git a/pylops/signalprocessing/fft.py b/pylops/signalprocessing/fft.py index 8322f75f..8343ea4d 100644 --- a/pylops/signalprocessing/fft.py +++ b/pylops/signalprocessing/fft.py @@ -306,7 +306,7 @@ def __init__( self.dims_t = dims_t # define padding(fftw requires the user to provide padded input signal) - self.pad = np.zeros((self.ndim, 2), dtype=int) + self.pad = np.zeros((self.ndims, 2), dtype=int) if self.real: if self.nfft % 2: self.pad[self.axis, 1] = ( diff --git a/pylops/signalprocessing/fft2d.py b/pylops/signalprocessing/fft2d.py index eb156159..12150184 100644 --- a/pylops/signalprocessing/fft2d.py +++ b/pylops/signalprocessing/fft2d.py @@ -56,7 +56,7 @@ def __init__( ) # checks - if self.ndim < 2: + if self.ndims < 2: msg = "FFT2D requires at least two input dimensions" raise ValueError(msg) if self.naxes != 2: @@ -174,7 +174,7 @@ def __init__( ) # checks - if self.ndim < 2: + if self.ndims < 2: msg = "FFT2D requires at least two input dimensions" raise ValueError(msg) if self.naxes != 2: @@ -284,7 +284,7 @@ def __init__( ) # checks - if self.ndim < 2: + if self.ndims < 2: msg = "FFT2D requires at least two input dimensions" raise ValueError(msg) if self.naxes != 2: diff --git a/pylops/utils/backend.py b/pylops/utils/backend.py index fd990d97..cdf52fa7 100644 --- a/pylops/utils/backend.py +++ b/pylops/utils/backend.py @@ -18,6 +18,9 @@ "get_sparse_eye", "get_lstsq", "get_sp_fft", + "get_eigvals", + "get_eigsh", + "get_lobpcg", "get_complex_dtype", "get_real_dtype", "to_cupy", @@ -36,10 +39,11 @@ import numpy as np import scipy.fft as sp_fft -from scipy.linalg import block_diag, lstsq, toeplitz +from scipy.linalg import block_diag, eigvals, lstsq, toeplitz from scipy.ndimage import gaussian_filter from scipy.signal import convolve, correlate, fftconvolve, oaconvolve from scipy.sparse import csc_matrix, csr_matrix, dia_matrix, eye +from scipy.sparse.linalg import eigsh, lobpcg from pylops.utils import deps from pylops.utils.typing import ArrayLike, DTypeLike, NDArray, Tfftengine_ncj @@ -48,6 +52,7 @@ import cupy as cp import cupyx import cupyx.scipy.fft as cp_fft + from cupy.linalg import eigvals as cp_eigvals from cupyx.scipy.linalg import block_diag as cp_block_diag from cupyx.scipy.linalg import toeplitz as cp_toeplitz from cupyx.scipy.ndimage import gaussian_filter as cp_gaussian_filter @@ -59,6 +64,8 @@ from cupyx.scipy.sparse import csr_matrix as cp_csr_matrix from cupyx.scipy.sparse import dia_matrix as cp_dia_matrix from cupyx.scipy.sparse import eye as cp_eye + from cupyx.scipy.sparse.linalg import eigsh as cp_eigsh + from cupyx.scipy.sparse.linalg import lobpcg as cp_lobpcg if deps.jax_enabled: import jax @@ -527,6 +534,75 @@ def get_sp_fft(x: ArrayLike) -> Callable: return cp_fft +def get_eigvals(x: ArrayLike) -> Callable: + """Returns correct eigvals module based on input + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Array + + Returns + ------- + f : :obj:`callable` + Function to be used to process array + + """ + if not deps.cupy_enabled: + return eigvals + + if cp.get_array_module(x) == np: + return eigvals + else: + return cp_eigvals + + +def get_eigsh(x: ArrayLike) -> Callable: + """Returns correct eigsh module based on input + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Array + + Returns + ------- + f : :obj:`callable` + Function to be used to process array + + """ + if not deps.cupy_enabled: + return eigsh + + if cp.get_array_module(x) == np: + return eigsh + else: + return cp_eigsh + + +def get_lobpcg(x: ArrayLike) -> Callable: + """Returns correct lobpcg module based on input + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Array + + Returns + ------- + f : :obj:`callable` + Function to be used to process array + + """ + if not deps.cupy_enabled: + return lobpcg + + if cp.get_array_module(x) == np: + return lobpcg + else: + return cp_lobpcg + + def get_complex_dtype(dtype: DTypeLike) -> DTypeLike: """Returns a complex type in the precision of the input type. diff --git a/pytests/test_linearoperator.py b/pytests/test_linearoperator.py index 238eb4ee..bf435b12 100644 --- a/pytests/test_linearoperator.py +++ b/pytests/test_linearoperator.py @@ -208,41 +208,288 @@ def test_sparse(par): assert_array_equal(S.toarray(), D) -@pytest.mark.skipif( - int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled" -) @pytest.mark.parametrize("par", [(par1), (par2), (par1j)]) -def test_eigs(par): - """Eigenvalues and condition number estimate with ARPACK""" +def test_eigs_lobpcg_failures(par): + """Fail uselobpcg when the operator is not symmetric or + when neigs is too large to avoid lobpcg dense solver + """ + nx = 10 * par["nx"] + + # lobpcg cannot be used for non-symmetric explicit operators (also when + # all eigenvalues are requested and the dense solver would be used) + Op = MatrixMult(np.diag(np.arange(nx, 0.0, -1.0)), dtype=np.dtype(par["dtype"])) + for neigs in (2, nx - 1, None): + with pytest.raises(ValueError, match="non real-symmetric or complex-hermitian"): + Op.eigs(neigs=neigs, symmetric=False, uselobpcg=True, backend=backend) + + # lobpcg cannot be used for non-symmetric linear operators + Op = Diagonal(np.arange(nx, 0, -1), dtype=par["dtype"]) + with pytest.raises(ValueError, match="non real-symmetric or complex-hermitian"): + Op.eigs(neigs=2, symmetric=False, uselobpcg=True, backend=backend) + + # lobpcg cannot be used when neigs is too large + with pytest.raises(ValueError, match="requires the operator size"): + Op.eigs(neigs=nx // 5 + 1, symmetric=True, uselobpcg=True, backend=backend) + # constraints reduce the available size + with pytest.raises(ValueError, match="number of constraints"): + Op.eigs( + neigs=nx // 5, + symmetric=True, + uselobpcg=True, + Y=np.eye(nx, 5, dtype=par["dtype"]), + backend=backend, + ) + + # lobpcg size check for non-square operators (based on number of columns) + Op = VStack([Op, Zero(par["ny"], nx)]) + with pytest.raises(ValueError, match="uselobpcg=True requires the operator size"): + Op.eigs(neigs=nx // 5 + 1, uselobpcg=True, backend=backend) + with pytest.raises(ValueError, match="uselobpcg=True requires the operator size"): + Op.eigs( + neigs=nx // 5, + uselobpcg=True, + Y=np.eye(nx, 5, dtype=par["dtype"]), + backend=backend, + ) + + +@pytest.mark.skipif(int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 0, reason="CuPy only") +@pytest.mark.parametrize("par", [(par1), (par1j)]) +def test_eigs_cupy_nonsymmetric_failures(par): + """Fail non-symmetric square operators with CuPy backend""" + nx = 10 * par["nx"] + diag = np.arange(nx, 0.0, -1.0) + + # explicit=True (limited number of eigenvalues) + Op = MatrixMult(np.diag(diag), dtype=np.dtype(par["dtype"])) + with pytest.raises(NotImplementedError, match="not available with CuPy"): + Op.eigs(neigs=3, symmetric=False, backend="cupy") + + # explicit=False + Op = Diagonal(diag, dtype=np.dtype(par["dtype"])) + with pytest.raises(NotImplementedError, match="not available with CuPy"): + Op.eigs(neigs=3, symmetric=False, backend="cupy") + + +@pytest.mark.parametrize("par", [(par1), (par1j)]) +@pytest.mark.parametrize("symmetric", [True, False]) +def test_eigs_square(par, symmetric): + """Eigenvalues and conditioning number estimate of square + diagonal matrix (note: by construction this is a symmetric + matrix with real-valued entries but not hermitian with + complex valued-entries) + """ + if symmetric and par["imag"] != 0: + pytest.skip( + "Skip complex symmetric case: cannot have complex-valued diagonal entries" + ) + if not symmetric and backend == "cupy": + pytest.skip("Skip non-symmetric case for CuPy as eigs is missing") + + # increase matrix size to ensure converge of lobpcg dense solver + nx = 10 * par["nx"] + # explicit=True - diag = np.arange(par["nx"], 0, -1) + par["imag"] * np.arange(par["nx"], 0, -1) + diag = np.arange(nx, 0.0, -1.0) + par["imag"] * np.arange(nx, 0.0, -1.0) + Op = MatrixMult(np.diag(diag), dtype=np.dtype(par["dtype"])) + + eigs = Op.eigs(symmetric=symmetric, backend=backend) # all + assert_array_almost_equal(eigs, diag, decimal=3) + + eigs = Op.eigs(neigs=3, symmetric=symmetric, backend=backend) # top-k + assert_array_almost_equal(eigs, diag[:3], decimal=3) + + # neigs equal to size of operator returns all eigenvalues + eigs = Op.eigs(neigs=nx, symmetric=symmetric, backend=backend) + assert_array_almost_equal(eigs, diag, decimal=3) + + if backend == "numpy": + # neigs equal to size of operator minus one (dense solver with + # selection based on which for symmetric=False, eigsh for + # symmetric=True; CuPy's eigsh does not support which="SM", + # which="BE", and sigma) + eigs = Op.eigs(neigs=nx - 1, symmetric=symmetric, backend=backend) + assert_array_almost_equal(eigs, diag[: nx - 1], decimal=3) + eigs = Op.eigs(neigs=nx - 1, symmetric=symmetric, which="SM", backend=backend) + assert_array_almost_equal(eigs, diag[1:], decimal=3) + + if symmetric: + eigs = Op.eigs(neigs=nx - 1, symmetric=True, sigma=0.0, backend=backend) + assert_array_almost_equal(eigs, diag[1:], decimal=3) + eigs = Op.eigs(neigs=nx - 1, symmetric=True, which="BE", backend=backend) + assert_array_almost_equal( + eigs, np.concatenate((diag[: nx // 2], diag[nx // 2 + 1 :])), decimal=3 + ) + + # CuPy cannot compute a single eigenvalue for square, non-symmetric matrix + cond = Op.cond(backend=backend) + assert_array_almost_equal(np.real(cond), nx, decimal=3) + + if symmetric: + # lobpcg + eigs = Op.eigs( + neigs=3, symmetric=True, uselobpcg=True, niter=50, backend=backend + ) + assert_array_almost_equal(eigs, diag[:3], decimal=3) + + # explicit=False + Op = Diagonal(diag, dtype=np.dtype(par["dtype"])) + + if backend == "numpy": + # CuPy cannot compute neigs-2 eigenvalues (seems to be + # a bug in CuPy's eigsh) + eigs = Op.eigs(symmetric=symmetric, backend=backend) + assert eigs.size == nx - 2 # neigs=None is clamped to N-2 (no warning) + assert_array_almost_equal(eigs, diag[: eigs.size], decimal=3) + + # neigs too large is clamped to N-1 for real symmetric operators + # (eigsh) and N-2 otherwise (eigs) (with warning) + neigsmax = nx - 1 if symmetric else nx - 2 + with pytest.warns(UserWarning, match="exceeds the number of columns"): + eigs = Op.eigs(neigs=nx, symmetric=symmetric, backend=backend) + assert eigs.size == neigsmax + assert_array_almost_equal(eigs, diag[:neigsmax], decimal=3) + + # neigs equal to N-1 is allowed for real symmetric operators (eigsh) + if symmetric: + eigs = Op.eigs(neigs=nx - 1, symmetric=True, backend=backend) + assert eigs.size == nx - 1 + assert_array_almost_equal(eigs, diag[: nx - 1], decimal=3) + + eigs = Op.eigs(neigs=3, symmetric=symmetric, backend=backend) # top-k + assert_array_almost_equal(eigs, diag[:3], decimal=3) + + if backend == "numpy": + cond = Op.cond(backend=backend) + assert_array_almost_equal(np.real(cond), nx, decimal=3) + + if symmetric: + # lobpcg + eigs1 = Op.eigs( + neigs=3, symmetric=True, uselobpcg=True, niter=50, backend=backend + ) + assert_array_almost_equal(eigs1, eigs, decimal=3) + + cond = Op.cond(uselobpcg=True, niter=100, backend=backend) + assert_array_almost_equal(np.real(cond), nx, decimal=3) + + +@pytest.mark.parametrize("par", [(par1j)]) +@pytest.mark.parametrize("explicit", [True, False]) +def test_eigs_square_hermitian(par, explicit): + """Eigenvalues and conditioning number estimate of complex-valued + hermitian matrix with known (real-valued) eigenvalues + """ + nx = 5 * par["nx"] + + # hermitian matrix A = U diag(eigs) U^H with U unitary + np.random.seed(0) + U, _ = np.linalg.qr( + np.random.normal(size=(nx, nx)) + 1j * np.random.normal(size=(nx, nx)) + ) + eigs_true = np.arange(nx, 0.0, -1.0) + A = (U * eigs_true) @ np.conj(U.T) + A = (A + np.conj(A.T)) / 2 # remove round-off asymmetry + Op = MatrixMult(A, dtype=np.dtype(par["dtype"])) + Op.explicit = explicit + + # hermitian solver + if explicit: + eigs = Op.eigs(symmetric=True, backend=backend) # all + assert_array_almost_equal(eigs, eigs_true, decimal=3) + + eigs = Op.eigs(neigs=3, symmetric=True, backend=backend) # top-k + assert_array_almost_equal(eigs, eigs_true[:3], decimal=3) + + if backend == "numpy": + cond = Op.cond(symmetric=True, backend=backend) + assert_array_almost_equal(np.real(cond), nx, decimal=3) + + # lobpcg + if backend == "numpy": + # CuPy cannot compute 2+ with lobpcg (seems to be + # a bug in CuPy's lobpcg) + eigs = Op.eigs( + neigs=3, symmetric=True, uselobpcg=True, niter=100, backend=backend + ) + assert_array_almost_equal(eigs, eigs_true[:3], decimal=3) + + # CuPy gives wrong answer (seems to be a bug in CuPy's lobpcg) + cond = Op.cond(uselobpcg=True, niter=200, backend=backend) + assert_array_almost_equal(np.real(cond), nx, decimal=3) + + # general solver (not available for CuPy) + if backend == "numpy": + eigs = Op.eigs(neigs=3, symmetric=False, backend=backend) + assert_array_almost_equal(np.real(eigs), eigs_true[:3], decimal=3) + assert_array_almost_equal(np.imag(eigs), np.zeros(3), decimal=3) + + +@pytest.mark.parametrize("par", [(par2), (par2j)]) +def test_eigs_non_square(par): + """Eigenvalues and conditioning number estimate of non-square matrix""" + + ny, nx = par["ny"], par["nx"] + + # explicit=True + diag = np.arange(nx, 0.0, -1.0) + par["imag"] * np.arange(nx, 0.0, -1.0) Op = MatrixMult( - np.vstack((np.diag(diag), np.zeros((par["ny"] - par["nx"], par["nx"])))) + np.vstack((np.diag(diag), np.zeros((ny - nx, nx)))), + dtype=np.dtype(par["dtype"]), ) - eigs = Op.eigs() - assert_array_almost_equal(diag[: eigs.size], eigs, decimal=3) - cond = Op.cond() - assert_array_almost_equal(np.real(cond), par["nx"], decimal=3) + eigs = Op.eigs(backend=backend) # all + assert_array_almost_equal(eigs, np.abs(diag)[: eigs.size], decimal=3) + + eigs = Op.eigs(neigs=3, backend=backend) # top-k + assert_array_almost_equal(eigs, np.abs(diag)[:3], decimal=3) + + # neigs equal to size of operator minus one + if backend == "numpy": + # CuPy cannot compute neigs-1 eigenvalues (seems to be + # a bug in CuPy's eigsh) + eigs = Op.eigs(neigs=nx - 1, backend=backend) + assert_array_almost_equal(eigs, np.abs(diag)[: nx - 1], decimal=3) + + # use lobpcg + if backend == "numpy": + # CuPy cannot compute 2+ with lobpcg (seems to be + # a bug in CuPy's lobpcg) + eigs = Op.eigs(neigs=2, uselobpcg=True, niter=100, backend=backend) + assert_array_almost_equal(eigs, np.abs(diag)[:2], decimal=3) + + if backend == "numpy": + cond = Op.cond(backend=backend) + assert_array_almost_equal(np.real(cond), nx, decimal=3) # explicit=False Op = Diagonal(diag, dtype=par["dtype"]) - if par["ny"] > par["nx"]: - Op = VStack([Op, Zero(par["ny"] - par["nx"], par["nx"])]) - eigs = Op.eigs() - assert_array_almost_equal(diag[: eigs.size], eigs, decimal=3) - - # uselobpcg cannot be used for square non-symmetric complex matrices - if np.iscomplex(Op): - eigs1 = Op.eigs(uselobpcg=True) - assert_array_almost_equal(eigs, eigs1, decimal=3) - - cond = Op.cond() - assert_array_almost_equal(np.real(cond), par["nx"], decimal=3) - - if np.iscomplex(Op): - cond1 = Op.cond(uselobpcg=True, niter=100) - assert_array_almost_equal(np.real(cond), np.real(cond1), decimal=3) + Op = VStack([Op, Zero(ny - nx, nx)]) + + if backend == "numpy": + # CuPy cannot compute neigs-2 eigenvalues (seems to be + # a bug in CuPy's eigsh) + eigs = Op.eigs(backend=backend) + assert_array_almost_equal(eigs, np.abs(diag)[: eigs.size], decimal=3) + + eigs = Op.eigs(neigs=3, backend=backend) # top-k + assert_array_almost_equal(eigs, np.abs(diag)[:3], decimal=3) + + if backend == "numpy": + cond = Op.cond(backend=backend) + assert_array_almost_equal(np.real(cond), nx, decimal=3) + + # use lobpcg + if backend == "numpy": + # CuPy cannot compute 2+ with lobpcg (seems to be + # a bug in CuPy's lobpcg) + eigs = Op.eigs(neigs=2, uselobpcg=True, niter=100, backend=backend) + assert_array_almost_equal(eigs, np.abs(diag)[:2], decimal=3) + + if backend == "numpy": + # CuPy gives wrong answer (seems to be a bug in CuPy's lobpcg) + cond = Op.cond(uselobpcg=True, niter=200, backend=backend) + assert_array_almost_equal(np.real(cond), nx, decimal=3) @pytest.mark.parametrize("par", [(par1), (par2), (par1j), (par2j)])