Merge pull request PyLops#612 from mrava87/patch-numpyinternal

Handle numpy warning for normalize_axis_index

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

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,23 +56,34 @@ 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
         self.weights = weights
         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
         )

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)
 

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

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:

pylops/signalprocessing/_baseffts.py

@@ -4,15 +4,18 @@
 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 (
     _raise_on_wrong_dtype,
     _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(

pylops/signalprocessing/convolvend.py

@@ -3,14 +3,14 @@
 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
 from pylops.utils.backend import (
     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)

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)

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