Feat: added sparse option to convmtx utility routines

pylops/avo/poststack.py

@@ -122,7 +122,9 @@ def _PoststackLinearModelling(
             )
         else:
             Cop = _MatrixMult(
-                nonstationary_convmtx(wav, nt0, hc=wav.shape[1] // 2, pad=(nt0, nt0)),
+                nonstationary_convmtx(
+                    wav, nt0, hc=wav.shape[1] // 2, pad=(nt0, nt0), sparse=sparse
+                ),
                 otherdims=spatdims,
                 dtype=dtype,
                 **args_MatrixMult,

pylops/basicoperators/matrixmult.py

@@ -23,7 +23,7 @@ class MatrixMult(LinearOperator):
 
     Parameters
     ----------
-    A : :obj:`numpy.ndarray` or :obj:`scipy.sparse` matrix
+    A : :obj:`numpy.ndarray` or :obj:`scipy.sparse.spmatrix`
         Matrix.
     otherdims : :obj:`tuple`, optional
         Number of samples for each other dimension of model

pylops/utils/backend.py

@@ -12,6 +12,8 @@
     "get_block_diag",
     "get_toeplitz",
     "get_csc_matrix",
+    "get_csr_matrix",
+    "get_dia_matrix",
     "get_sparse_eye",
     "get_lstsq",
     "get_sp_fft",
@@ -35,7 +37,7 @@
 import scipy.fft as sp_fft
 from scipy.linalg import block_diag, lstsq, toeplitz
 from scipy.signal import convolve, correlate, fftconvolve, oaconvolve
-from scipy.sparse import csc_matrix, eye
+from scipy.sparse import csc_matrix, csr_matrix, dia_matrix, eye
 
 from pylops.utils import deps
 from pylops.utils.typing import ArrayLike, DTypeLike, NDArray, Tfftengine_ncj
@@ -51,6 +53,8 @@
     from cupyx.scipy.signal import fftconvolve as cp_fftconvolve
     from cupyx.scipy.signal import oaconvolve as cp_oaconvolve
     from cupyx.scipy.sparse import csc_matrix as cp_csc_matrix
+    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
 
 if deps.jax_enabled:
@@ -382,6 +386,52 @@ def get_csc_matrix(x: ArrayLike) -> Callable:
         return cp_csc_matrix
 
 
+def get_csr_matrix(x: ArrayLike) -> Callable:
+    """Returns correct csr_matrix 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 csr_matrix
+
+    if cp.get_array_module(x) == np:
+        return csr_matrix
+    else:
+        return cp_csr_matrix
+
+
+def get_dia_matrix(x: ArrayLike) -> Callable:
+    """Returns correct dia_matrix 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 dia_matrix
+
+    if cp.get_array_module(x) == np:
+        return dia_matrix
+    else:
+        return cp_dia_matrix
+
+
 def get_sparse_eye(x: ArrayLike) -> Callable:
     """Returns correct sparse eye module based on input
 

pylops/utils/signalprocessing.py

@@ -18,13 +18,15 @@
 from pylops.utils._pwd2d import _conv_allpass, _triangular_smoothing_from_boxcars
 from pylops.utils.backend import (
     get_array_module,
+    get_csr_matrix,
+    get_dia_matrix,
     get_normalize_axis_index,
     get_toeplitz,
 )
 from pylops.utils.typing import NDArray, Tpwdsmoothing
 
 
-def convmtx(h: NDArray, n: int, offset: int = 0) -> NDArray:
+def convmtx(h: NDArray, n: int, offset: int = 0, sparse: bool = False) -> NDArray:
     r"""Convolution matrix
 
     Makes a dense convolution matrix :math:`\mathbf{C}`
@@ -42,12 +44,16 @@ def convmtx(h: NDArray, n: int, offset: int = 0) -> NDArray:
         Convolution filter (1D array)
     n : :obj:`int`
         Number of columns of convolution matrix
-    offset : :obj:`int`
+    offset : :obj:`int`, optional
         Index of the center of the filter
+    sparse : :obj:`bool`, optional
+        .. versionadded:: 2.9.0
+
+        Return dense (``False``) or sparse (``True``) matrix
 
     Returns
     -------
-    C : :obj:`numpy.ndarray`
+    C : :obj:`numpy.ndarray` or :obj:`scipy.sparse.spmatrix`
         Convolution matrix of size :math:`\text{len}(h)+n-1 \times n`
 
     """
@@ -62,12 +68,26 @@ def convmtx(h: NDArray, n: int, offset: int = 0) -> NDArray:
     )
 
     ncp = get_array_module(h)
+
+    # create Toeplitz matrix
     nh = len(h)
     col_1 = ncp.r_[h, ncp.zeros(n + nh - 2, dtype=h.dtype)]
     row_1 = ncp.r_[h[0], ncp.zeros(n - 1, dtype=h.dtype)]
     C = get_toeplitz(h)(col_1, row_1)
+
     # apply offset
     C = C[offset : offset + nh + n - 1]
+
+    # convert to sparse using the following rule-of-thumb:
+    # - DIA format very short filters (<= 11)
+    # - CSR format for other filters
+    if sparse:
+        if ncp == np:
+            C = get_dia_matrix(h)(C) if nh <= 11 else get_csr_matrix(h)(C)
+        else:
+            # For CuPy DIA cannot take a dense matrix, so the dense matrix is
+            # always converted to CSR format
+            C = get_csr_matrix(h)(C)
     return C
 
 
@@ -76,6 +96,7 @@ def nonstationary_convmtx(
     n: int,
     hc: int = 0,
     pad: tuple[int] = (0, 0),
+    sparse: bool = False,
 ) -> NDArray:
     r"""Convolution matrix from a bank of filters
 
@@ -97,18 +118,35 @@ def nonstationary_convmtx(
         Zero-padding to apply to the bank of filters before and after the
         provided values (use it to avoid wrap-around or pass filters with
         enough padding)
+    sparse : :obj:`bool`, optional
+        .. versionadded:: 2.9.0
+
+        Return dense (``False``) or sparse (``True``) matrix
 
     Returns
     -------
-    C : :obj:`numpy.ndarray`
+    C : :obj:`numpy.ndarray` or :obj:`scipy.sparse.spmatrix`
         Convolution matrix
 
     """
     ncp = get_array_module(H)
 
+    # create Toeplitz matrix
+    nh = H.shape[1]
     H = ncp.pad(H, ((0, 0), pad), mode="constant")
     C = ncp.array([ncp.roll(h, ih) for ih, h in enumerate(H)])
     C = C[:, pad[0] + hc : pad[0] + hc + n].T  # take away edges
+
+    # convert to sparse using the following rule-of-thumb:
+    # - DIA format very short filters (<= 11)
+    # - CSR format for other filters
+    if sparse:
+        if ncp == np:
+            C = get_dia_matrix(H)(C) if nh <= 11 else get_csr_matrix(H)(C)
+        else:
+            # For CuPy DIA cannot take a dense matrix, so the dense matrix is
+            # always converted to CSR format
+            C = get_csr_matrix(H)(C)
     return C
 
 

pytests/test_signalutils.py

@@ -1,8 +1,16 @@
 import os
 
-import numpy as np
+if int(os.environ.get("TEST_CUPY_PYLOPS", 0)):
+    import cupy as np
+    from cupy.testing import assert_array_almost_equal
+
+    backend = "cupy"
+else:
+    import numpy as np
+    from numpy.testing import assert_array_almost_equal
+
+    backend = "numpy"
 import pytest
-from numpy.testing import assert_array_almost_equal
 
 from pylops.utils.signalprocessing import convmtx, nonstationary_convmtx, slope_estimate
 
@@ -39,57 +47,57 @@
 np.random.seed(10)
 
 
-@pytest.mark.skipif(
-    int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled"
-)
 @pytest.mark.parametrize("par", [(par1), (par1j), (par2), (par2j)])
-def test_convmtx(par):
+@pytest.mark.parametrize("sparse", [False, True])
+def test_convmtx(par, sparse):
     """Compare convmtx with np.convolve (small filter)"""
     x = np.random.normal(0, 1, par["nt"]) + par["imag"] * np.random.normal(
         0, 1, par["nt"]
     )
 
     h = np.hanning(par["nh"])
-    H = convmtx(h, par["nt"], par["nh"] // 2)
+    H = convmtx(h, par["nt"], par["nh"] // 2, sparse=sparse)
 
     y = np.convolve(x, h, mode="same")
-    y1 = np.dot(H[: par["nt"]], x)
+    y1 = (H @ x)[: par["nt"]]
     assert_array_almost_equal(y, y1, decimal=4)
 
 
-@pytest.mark.skipif(
-    int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled"
-)
 @pytest.mark.parametrize("par", [(par1), (par1j), (par2), (par2j)])
-def test_convmtx1(par):
+@pytest.mark.parametrize("sparse", [False, True])
+def test_convmtx1(par, sparse):
     """Compare convmtx with np.convolve (large filter)"""
     x = np.random.normal(0, 1, par["nt"]) + par["imag"] * np.random.normal(
         0, 1, par["nt"]
     )
 
     h = np.hanning(par["nh"])
     X = convmtx(
-        x, par["nh"], par["nh"] // 2 - 1 if par["nh"] % 2 == 0 else par["nh"] // 2
+        x,
+        par["nh"],
+        par["nh"] // 2 - 1 if par["nh"] % 2 == 0 else par["nh"] // 2,
+        sparse=sparse,
     )
 
     y = np.convolve(x, h, mode="same")
-    y1 = np.dot(X[: par["nt"]], h)
+    y1 = (X @ h)[: par["nt"]]
     assert_array_almost_equal(y, y1, decimal=4)
 
 
-@pytest.mark.skipif(
-    int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled"
-)
 @pytest.mark.parametrize("par", [(par1), (par1j)])
-def test_nonstationary_convmtx(par):
+@pytest.mark.parametrize("sparse", [False, True])
+def test_nonstationary_convmtx(par, sparse):
     """Compare nonstationary_convmtx with convmtx for stationary filter"""
     x = np.random.normal(0, 1, par["nt"]) + par["imag"] * np.random.normal(
         0, 1, par["nt"]
     )
 
     h = np.hanning(par["nh"])
     H = convmtx(
-        h, par["nt"], par["nh"] // 2 - 1 if par["nh"] % 2 == 0 else par["nh"] // 2
+        h,
+        par["nt"],
+        par["nh"] // 2 - 1 if par["nh"] % 2 == 0 else par["nh"] // 2,
+        sparse=sparse,
     )
 
     H1 = nonstationary_convmtx(
@@ -98,7 +106,8 @@ def test_nonstationary_convmtx(par):
         hc=par["nh"] // 2,
         pad=(par["nt"], par["nt"]),
     )
-    y = np.dot(H[: par["nt"]], x)
+
+    y = (H @ x)[: par["nt"]]
     y1 = np.dot(H1, x)
     assert_array_almost_equal(y, y1, decimal=4)