Merge pull request #772 from mrava87/feat-slopeest_cupy

Feat: slope_estimate cupy backend

Author: mrava87

pylops/utils/backend.py

@@ -11,6 +11,7 @@
     "get_sliding_window_view",
     "get_block_diag",
     "get_toeplitz",
+    "get_gaussian_filter",
     "get_csc_matrix",
     "get_csr_matrix",
     "get_dia_matrix",
@@ -36,6 +37,7 @@
 import numpy as np
 import scipy.fft as sp_fft
 from scipy.linalg import block_diag, 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
 
@@ -48,6 +50,7 @@
     import cupyx.scipy.fft as cp_fft
     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
     from cupyx.scipy.signal import convolve as cp_convolve
     from cupyx.scipy.signal import correlate as cp_correlate
     from cupyx.scipy.signal import fftconvolve as cp_fftconvolve
@@ -363,6 +366,29 @@ def get_toeplitz(x: ArrayLike) -> Callable:
         return toeplitz
 
 
+def get_gaussian_filter(x: ArrayLike) -> Callable:
+    """Returns correct gaussian_filter 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 gaussian_filter
+
+    if cp.get_array_module(x) == np:
+        return gaussian_filter
+    else:
+        return cp_gaussian_filter
+
+
 def get_csc_matrix(x: ArrayLike) -> Callable:
     """Returns correct csc_matrix module based on input
 

pylops/utils/signalprocessing.py

@@ -10,7 +10,6 @@
 from collections.abc import Sequence
 
 import numpy as np
-from scipy.ndimage import gaussian_filter
 
 from pylops.basicoperators import Diagonal, Smoothing2D, SmoothingND
 from pylops.optimization.leastsquares import preconditioned_inversion
@@ -20,6 +19,7 @@
     get_array_module,
     get_csr_matrix,
     get_dia_matrix,
+    get_gaussian_filter,
     get_normalize_axis_index,
     get_toeplitz,
 )
@@ -266,37 +266,39 @@ def slope_estimate(
         anisotropy in digitized images", Journal ASCI Imaging Workshop. 1995.
 
     """
-    slopes = np.zeros_like(d)
-    anisos = np.zeros_like(d)
+    ncp = get_array_module(d)
 
-    gz, gx = np.gradient(d, dz, dx)
+    slopes = ncp.zeros_like(d)
+    anisos = ncp.zeros_like(d)
+
+    gz, gx = ncp.gradient(d, dz, dx)
     gzz, gzx, gxx = gz * gz, gz * gx, gx * gx
 
     # smoothing
-    gzz = gaussian_filter(gzz, sigma=smooth)
-    gzx = gaussian_filter(gzx, sigma=smooth)
-    gxx = gaussian_filter(gxx, sigma=smooth)
+    gzz = get_gaussian_filter(d)(gzz, sigma=smooth)
+    gzx = get_gaussian_filter(d)(gzx, sigma=smooth)
+    gxx = get_gaussian_filter(d)(gxx, sigma=smooth)
 
-    gmax = max(gzz.max(), gxx.max(), np.abs(gzx).max())
+    gmax = max(gzz.max(), gxx.max(), ncp.abs(gzx).max())
     if gmax <= eps:
-        return np.zeros_like(d), anisos
+        return ncp.zeros_like(d), anisos
 
     gzz /= gmax
     gzx /= gmax
     gxx /= gmax
 
     lcommon1 = 0.5 * (gzz + gxx)
-    lcommon2 = 0.5 * np.sqrt((gzz - gxx) ** 2 + 4 * gzx**2)
+    lcommon2 = 0.5 * ncp.sqrt((gzz - gxx) ** 2 + 4 * gzx**2)
     l1 = lcommon1 + lcommon2
     l2 = lcommon1 - lcommon2
 
     regdata = l1 > eps
     anisos[regdata] = 1 - l2[regdata] / l1[regdata]
 
     if dips:
-        slopes = 0.5 * np.arctan2(2 * gzx, gzz - gxx)
+        slopes = 0.5 * ncp.arctan2(2 * gzx, gzz - gxx)
     else:
-        regdata = np.abs(gzx) > eps
+        regdata = ncp.abs(gzx) > eps
         slopes[regdata] = (l1 - gzz)[regdata] / gzx[regdata]
 
     return slopes, anisos

pytests/test_signalutils.py

@@ -112,9 +112,6 @@ def test_nonstationary_convmtx(par, sparse):
     assert_array_almost_equal(y, y1, decimal=4)
 
 
-@pytest.mark.skipif(
-    int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled"
-)
 def test_slope_estimation_dips():
     """Slope estimation using the Structure tensor algorithm should
     apply regularisation (some slopes are set to zero)