diff --git a/pylops/basicoperators/firstderivative.py b/pylops/basicoperators/firstderivative.py index f8bd208e..c19a2650 100644 --- a/pylops/basicoperators/firstderivative.py +++ b/pylops/basicoperators/firstderivative.py @@ -3,11 +3,15 @@ from typing import Callable, Union import numpy as np -from numpy.core.multiarray import normalize_axis_index from pylops import LinearOperator from pylops.utils._internal import _value_or_sized_to_tuple -from pylops.utils.backend import get_array_module, inplace_add, inplace_set +from pylops.utils.backend import ( + get_array_module, + get_normalize_axis_index, + inplace_add, + inplace_set, +) from pylops.utils.decorators import reshaped from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray @@ -95,7 +99,7 @@ def __init__( dims = _value_or_sized_to_tuple(dims) super().__init__(dtype=np.dtype(dtype), dims=dims, dimsd=dims, name=name) - self.axis = normalize_axis_index(axis, len(self.dims)) + self.axis = get_normalize_axis_index()(axis, len(self.dims)) self.sampling = sampling self.kind = kind self.edge = edge diff --git a/pylops/basicoperators/laplacian.py b/pylops/basicoperators/laplacian.py index d2f3e64c..f8f836a5 100644 --- a/pylops/basicoperators/laplacian.py +++ b/pylops/basicoperators/laplacian.py @@ -2,13 +2,11 @@ from typing import Tuple -from pylops.utils.typing import NDArray - -from numpy.core.multiarray import normalize_axis_index from pylops import LinearOperator from pylops.basicoperators import SecondDerivative -from pylops.utils.typing import DTypeLike, InputDimsLike +from pylops.utils.backend import get_normalize_axis_index +from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray class Laplacian(LinearOperator): @@ -58,14 +56,18 @@ class Laplacian(LinearOperator): """ - def __init__(self, dims: InputDimsLike, - axes: InputDimsLike = (-2, -1), - weights: Tuple[float, ...] = (1, 1), - sampling: Tuple[float, ...] = (1, 1), - edge: bool = False, - kind: str = "centered", - dtype: DTypeLike = "float64", name: str = "L"): - axes = tuple(normalize_axis_index(ax, len(dims)) for ax in axes) + def __init__( + self, + dims: InputDimsLike, + axes: InputDimsLike = (-2, -1), + weights: Tuple[float, ...] = (1, 1), + sampling: Tuple[float, ...] = (1, 1), + edge: bool = False, + kind: str = "centered", + dtype: DTypeLike = "float64", + name: str = "L", + ): + axes = tuple(get_normalize_axis_index()(ax, len(dims)) for ax in axes) if not (len(axes) == len(weights) == len(sampling)): raise ValueError("axes, weights, and sampling have different size") self.axes = axes @@ -73,8 +75,15 @@ def __init__(self, dims: InputDimsLike, self.sampling = sampling self.edge = edge self.kind = kind - Op = self._calc_l2op(dims=dims, axes=axes, sampling=sampling, edge=edge, kind=kind, dtype=dtype, - weights=weights) + Op = self._calc_l2op( + dims=dims, + axes=axes, + sampling=sampling, + edge=edge, + kind=kind, + dtype=dtype, + weights=weights, + ) super().__init__(Op=Op, name=name) def _matvec(self, x: NDArray) -> NDArray: @@ -84,8 +93,15 @@ def _rmatvec(self, x: NDArray) -> NDArray: return super()._rmatvec(x) @staticmethod - def _calc_l2op(dims: InputDimsLike, axes: InputDimsLike, weights: Tuple[float, ...], sampling: Tuple[float, ...], - edge: bool, kind: str, dtype: DTypeLike): + def _calc_l2op( + dims: InputDimsLike, + axes: InputDimsLike, + weights: Tuple[float, ...], + sampling: Tuple[float, ...], + edge: bool, + kind: str, + dtype: DTypeLike, + ): l2op = SecondDerivative( dims, axis=axes[0], sampling=sampling[0], edge=edge, kind=kind, dtype=dtype ) diff --git a/pylops/basicoperators/restriction.py b/pylops/basicoperators/restriction.py index 1a745b30..2595c971 100644 --- a/pylops/basicoperators/restriction.py +++ b/pylops/basicoperators/restriction.py @@ -6,17 +6,14 @@ import numpy as np import numpy.ma as np_ma -# need to check numpy version since normalize_axis_index will be -# soon moved from numpy.core.multiarray to from numpy.lib.array_utils -np_version = np.__version__.split(".") -if int(np_version[0]) < 2: - from numpy.core.multiarray import normalize_axis_index -else: - from numpy.lib.array_utils import normalize_axis_index - from pylops import LinearOperator from pylops.utils._internal import _value_or_sized_to_tuple -from pylops.utils.backend import get_array_module, inplace_set, to_cupy_conditional +from pylops.utils.backend import ( + get_array_module, + get_normalize_axis_index, + inplace_set, + to_cupy_conditional, +) from pylops.utils.typing import DTypeLike, InputDimsLike, IntNDArray, NDArray logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.WARNING) @@ -119,7 +116,7 @@ def __init__( ) -> None: ncp = get_array_module(iava) dims = _value_or_sized_to_tuple(dims) - axis = normalize_axis_index(axis, len(dims)) + axis = get_normalize_axis_index()(axis, len(dims)) dimsd = list(dims) # data dimensions dimsd[axis] = len(iava) diff --git a/pylops/basicoperators/secondderivative.py b/pylops/basicoperators/secondderivative.py index 8433987d..d7fe7b8b 100644 --- a/pylops/basicoperators/secondderivative.py +++ b/pylops/basicoperators/secondderivative.py @@ -3,11 +3,15 @@ from typing import Callable, Union import numpy as np -from numpy.core.multiarray import normalize_axis_index from pylops import LinearOperator from pylops.utils._internal import _value_or_sized_to_tuple -from pylops.utils.backend import get_array_module, inplace_add, inplace_set +from pylops.utils.backend import ( + get_array_module, + get_normalize_axis_index, + inplace_add, + inplace_set, +) from pylops.utils.decorators import reshaped from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray @@ -86,7 +90,7 @@ def __init__( dims = _value_or_sized_to_tuple(dims) super().__init__(dtype=np.dtype(dtype), dims=dims, dimsd=dims, name=name) - self.axis = normalize_axis_index(axis, len(self.dims)) + self.axis = get_normalize_axis_index()(axis, len(self.dims)) self.sampling = sampling self.kind = kind self.edge = edge diff --git a/pylops/basicoperators/transpose.py b/pylops/basicoperators/transpose.py index 0ccf05f0..88c99514 100644 --- a/pylops/basicoperators/transpose.py +++ b/pylops/basicoperators/transpose.py @@ -1,10 +1,10 @@ __all__ = ["Transpose"] import numpy as np -from numpy.core.multiarray import normalize_axis_index from pylops import LinearOperator from pylops.utils._internal import _value_or_sized_to_tuple +from pylops.utils.backend import get_normalize_axis_index from pylops.utils.decorators import reshaped from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray @@ -64,7 +64,7 @@ def __init__( ) -> None: dims = _value_or_sized_to_tuple(dims) ndims = len(dims) - self.axes = [normalize_axis_index(ax, ndims) for ax in axes] + self.axes = [get_normalize_axis_index()(ax, ndims) for ax in axes] # find out if all axes are present only once in axes if len(np.unique(self.axes)) != ndims: diff --git a/pylops/signalprocessing/_baseffts.py b/pylops/signalprocessing/_baseffts.py index cb8c1b41..18a4cce2 100644 --- a/pylops/signalprocessing/_baseffts.py +++ b/pylops/signalprocessing/_baseffts.py @@ -4,7 +4,6 @@ from typing import Optional, Sequence, Union import numpy as np -from numpy.core.multiarray import normalize_axis_index from pylops import LinearOperator from pylops.utils._internal import ( @@ -12,7 +11,11 @@ _value_or_sized_to_array, _value_or_sized_to_tuple, ) -from pylops.utils.backend import get_complex_dtype, get_real_dtype +from pylops.utils.backend import ( + get_complex_dtype, + get_normalize_axis_index, + get_real_dtype, +) from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.WARNING) @@ -46,7 +49,7 @@ def __init__( axes = _value_or_sized_to_array(axis) _raise_on_wrong_dtype(axes, np.integer, "axis") - self.axis = normalize_axis_index(axes[0], self.ndim) + self.axis = get_normalize_axis_index()(axes[0], self.ndim) if nfft is None: self.nfft = dims[self.axis] @@ -152,7 +155,7 @@ def __init__( axes = _value_or_sized_to_array(axes) _raise_on_wrong_dtype(axes, np.integer, "axes") - self.axes = np.array([normalize_axis_index(d, self.ndim) for d in axes]) + self.axes = np.array([get_normalize_axis_index()(d, self.ndim) for d in axes]) self.naxes = len(self.axes) if self.naxes != len(np.unique(self.axes)): warnings.warn( diff --git a/pylops/signalprocessing/convolvend.py b/pylops/signalprocessing/convolvend.py index c2fcb563..14ea6e24 100644 --- a/pylops/signalprocessing/convolvend.py +++ b/pylops/signalprocessing/convolvend.py @@ -3,7 +3,6 @@ from typing import Optional, Union import numpy as np -from numpy.core.multiarray import normalize_axis_index from pylops import LinearOperator from pylops.utils._internal import _value_or_sized_to_tuple @@ -11,6 +10,7 @@ get_array_module, get_convolve, get_correlate, + get_normalize_axis_index, to_cupy_conditional, ) from pylops.utils.decorators import reshaped @@ -78,7 +78,9 @@ def __init__( self.axes = ( np.arange(len(self.dims)) if axes is None - else np.array([normalize_axis_index(ax, len(self.dims)) for ax in axes]) + else np.array( + [get_normalize_axis_index()(ax, len(self.dims)) for ax in axes] + ) ) self.h = h hshape = np.array(self.h.shape) diff --git a/pylops/signalprocessing/shift.py b/pylops/signalprocessing/shift.py index b7d17ed7..c94b878e 100644 --- a/pylops/signalprocessing/shift.py +++ b/pylops/signalprocessing/shift.py @@ -4,11 +4,11 @@ import numpy as np import numpy.typing as npt -from numpy.core.multiarray import normalize_axis_index from pylops.basicoperators import Diagonal from pylops.signalprocessing import FFT from pylops.utils._internal import _value_or_sized_to_array +from pylops.utils.backend import get_normalize_axis_index from pylops.utils.typing import DTypeLike @@ -113,7 +113,7 @@ def Shift( Sop = Diagonal(shift, dims=dimsdiag, axis=axis, dtype=Fop.cdtype) else: # add dimensions to shift to match dimensions of model and data - axis = normalize_axis_index(axis, len(dims)) + axis = get_normalize_axis_index()(axis, len(dims)) fdims = np.ones(shift.ndim + 1, dtype=int) fdims[axis] = Fop.f.size f = Fop.f.reshape(fdims) diff --git a/pylops/utils/backend.py b/pylops/utils/backend.py index 50da7faa..9b7d3b2d 100644 --- a/pylops/utils/backend.py +++ b/pylops/utils/backend.py @@ -2,6 +2,7 @@ "get_module", "get_module_name", "get_array_module", + "get_normalize_axis_index", "get_convolve", "get_fftconvolve", "get_oaconvolve", @@ -59,6 +60,14 @@ from jax.scipy.signal import convolve as j_convolve from jax.scipy.signal import fftconvolve as j_fftconvolve +# need to check numpy version since the namespace of normalize_axis_index +# changed from numpy>=2.0.0 +np_version = np.__version__.split(".") +if int(np_version[0]) > 1: + from numpy.lib.array_utils import normalize_axis_index +else: + from numpy.core.multiarray import normalize_axis_index + def get_module(backend: str = "numpy") -> ModuleType: """Returns correct numerical module based on backend string @@ -138,6 +147,18 @@ def get_array_module(x: npt.ArrayLike) -> ModuleType: return np +def get_normalize_axis_index() -> Callable: + """Returns correct normalize_axis_index module based on numpy version + + Returns + ------- + f : :obj:`func` + Function to be used to process array + + """ + return normalize_axis_index + + def get_convolve(x: npt.ArrayLike) -> Callable: """Returns correct convolve module based on input