Merge pull request #816 from mrava87/feat-eigs_cupy

Feat: added CuPy support to LinearOperator.eigs and LinearOperator.cond

Author: mrava87

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``.

pylops/linearoperator.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",

pylops/optimization/eigs.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)
             ):

pylops/signalprocessing/_baseffts.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] = (

pylops/signalprocessing/dtcwt.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] = (

pylops/signalprocessing/fft.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:

pylops/signalprocessing/fft2d.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.