diff --git a/pylops/utils/backend.py b/pylops/utils/backend.py index 8347d71d..fd990d97 100644 --- a/pylops/utils/backend.py +++ b/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 diff --git a/pylops/utils/signalprocessing.py b/pylops/utils/signalprocessing.py index faeabd9b..2a0044e6 100644 --- a/pylops/utils/signalprocessing.py +++ b/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,27 +266,29 @@ 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 @@ -294,9 +296,9 @@ def slope_estimate( 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 diff --git a/pytests/test_signalutils.py b/pytests/test_signalutils.py index d1857f8e..5feffbbf 100644 --- a/pytests/test_signalutils.py +++ b/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)