raise warning similar to numpy >= 2.0 (#139)

* raise warning similar to numpy

* changes to raise warning only when numpy > 2.0 is used

* remove duplication

* address comments

Author: vtavana

README.md

@@ -82,5 +82,5 @@ The package also provides `mkl_fft._numpy_fft` and `mkl_fft._scipy_fft` interfac
 
 To build ``mkl_fft`` from sources on Linux:
   - install a recent version of MKL, if necessary;
-  - execute ``source /path_to_oneapi/mkl/latest/env/vars.sh`` ;
+  - execute ``source /path_to_oneapi/mkl/latest/env/vars.sh``;
   - execute ``python -m pip install .``

mkl_fft/_float_utils.py

@@ -81,7 +81,7 @@ def __downcast_float128_array(x):
 
 def __supported_array_or_not_implemented(x):
     """
-    Used in _scipy_fft_backend to convert array to float32,
+    Used in _scipy_fft to convert array to float32,
     float64, complex64, or complex128 type or return NotImplemented
     """
     __x = np.asarray(x)

mkl_fft/_numpy_fft.py

@@ -71,36 +71,47 @@
 ]
 
 import re
+import warnings
 
-from numpy import array, asanyarray, conjugate, prod, sqrt, take
+import numpy as np
+from numpy import array, conjugate, prod, sqrt, take
 
 from . import _float_utils
 from . import _pydfti as mkl_fft  # pylint: disable=no-name-in-module
 
 
+def _compute_fwd_scale(norm, n, shape):
+    _check_norm(norm)
+    if norm in (None, "backward"):
+        return 1.0
+
+    ss = n if n is not None else shape
+    nn = prod(ss)
+    fsc = 1 / nn if nn != 0 else 1
+    if norm == "forward":
+        return fsc
+    else:  # norm == "ortho"
+        return sqrt(fsc)
+
+
 def _check_norm(norm):
     if norm not in (None, "ortho", "forward", "backward"):
         raise ValueError(
-            (
-                "Invalid norm value {} should be None, "
-                '"ortho", "forward", or "backward".'
-            ).format(norm)
+            f"Invalid norm value {norm} should be None, 'ortho', 'forward', "
+            "or 'backward'."
         )
 
 
-def frwd_sc_1d(n, s):
-    nn = n if n is not None else s
-    return 1 / nn if nn != 0 else 1
-
-
-def frwd_sc_nd(s, x_shape):
-    ss = s if s is not None else x_shape
-    nn = prod(ss)
-    return 1 / nn if nn != 0 else 1
-
+def _swap_direction(norm):
+    _check_norm(norm)
+    _swap_direction_map = {
+        "backward": "forward",
+        None: "forward",
+        "ortho": "ortho",
+        "forward": "backward",
+    }
 
-def ortho_sc_1d(n, s):
-    return sqrt(frwd_sc_1d(n, s))
+    return _swap_direction_map[norm]
 
 
 def trycall(func, args, kwrds):
@@ -206,15 +217,9 @@ def fft(a, n=None, axis=-1, norm=None):
     the `numpy.fft` documentation.
 
     """
-    _check_norm(norm)
-    x = _float_utils.__downcast_float128_array(a)
 
-    if norm in (None, "backward"):
-        fsc = 1.0
-    elif norm == "forward":
-        fsc = frwd_sc_1d(n, x.shape[axis])
-    else:
-        fsc = ortho_sc_1d(n, x.shape[axis])
+    x = _float_utils.__downcast_float128_array(a)
+    fsc = _compute_fwd_scale(norm, n, x.shape[axis])
 
     return trycall(mkl_fft.fft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc})
 
@@ -305,15 +310,9 @@ def ifft(a, n=None, axis=-1, norm=None):
     >>> plt.show()
 
     """
-    _check_norm(norm)
-    x = _float_utils.__downcast_float128_array(a)
 
-    if norm in (None, "backward"):
-        fsc = 1.0
-    elif norm == "forward":
-        fsc = frwd_sc_1d(n, x.shape[axis])
-    else:
-        fsc = ortho_sc_1d(n, x.shape[axis])
+    x = _float_utils.__downcast_float128_array(a)
+    fsc = _compute_fwd_scale(norm, n, x.shape[axis])
 
     return trycall(mkl_fft.ifft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc})
 
@@ -402,15 +401,9 @@ def rfft(a, n=None, axis=-1, norm=None):
     exploited to compute only the non-negative frequency terms.
 
     """
-    _check_norm(norm)
-    x = _float_utils.__downcast_float128_array(a)
 
-    if norm in (None, "backward"):
-        fsc = 1.0
-    elif norm == "forward":
-        fsc = frwd_sc_1d(n, x.shape[axis])
-    else:
-        fsc = ortho_sc_1d(n, x.shape[axis])
+    x = _float_utils.__downcast_float128_array(a)
+    fsc = _compute_fwd_scale(norm, n, x.shape[axis])
 
     return trycall(mkl_fft.rfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc})
 
@@ -501,16 +494,9 @@ def irfft(a, n=None, axis=-1, norm=None):
     specified, and the output array is purely real.
 
     """
-    _check_norm(norm)
-    x = _float_utils.__downcast_float128_array(a)
 
-    nn = n if n else 2 * (x.shape[axis] - 1)
-    if norm in (None, "backward"):
-        fsc = 1.0
-    elif norm == "forward":
-        fsc = frwd_sc_1d(nn, nn)
-    else:
-        fsc = ortho_sc_1d(nn, nn)
+    x = _float_utils.__downcast_float128_array(a)
+    fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1))
 
     return trycall(
         mkl_fft.irfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc}
@@ -593,18 +579,12 @@ def hfft(a, n=None, axis=-1, norm=None):
            [ 2., -2.]])
 
     """
-    _check_norm(norm)
+
+    norm = _swap_direction(norm)
     x = _float_utils.__downcast_float128_array(a)
     x = array(x, copy=True, dtype=complex)
     conjugate(x, out=x)
-
-    nn = n if n else 2 * (x.shape[axis] - 1)
-    if norm in (None, "backward"):
-        fsc = frwd_sc_1d(nn, nn)
-    elif norm == "forward":
-        fsc = 1.0
-    else:
-        fsc = ortho_sc_1d(nn, nn)
+    fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1))
 
     return trycall(
         mkl_fft.irfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc}
@@ -668,17 +648,12 @@ def ihfft(a, n=None, axis=-1, norm=None):
     array([ 1.-0.j,  2.-0.j,  3.-0.j,  4.-0.j])
 
     """
+
     # The copy may be required for multithreading.
-    _check_norm(norm)
+    norm = _swap_direction(norm)
     x = _float_utils.__downcast_float128_array(a)
     x = array(x, copy=True, dtype=float)
-
-    if norm in (None, "backward"):
-        fsc = frwd_sc_1d(n, x.shape[axis])
-    elif norm == "forward":
-        fsc = 1.0
-    else:
-        fsc = ortho_sc_1d(n, x.shape[axis])
+    fsc = _compute_fwd_scale(norm, n, x.shape[axis])
 
     output = trycall(
         mkl_fft.rfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc}
@@ -688,22 +663,46 @@ def ihfft(a, n=None, axis=-1, norm=None):
     return output
 
 
-def _cook_nd_args(a, s=None, axes=None, invreal=0):
+# copied from: https://github.com/numpy/numpy/blob/main/numpy/fft/_pocketfft.py
+def _cook_nd_args(a, s=None, axes=None, invreal=False):
     if s is None:
-        shapeless = 1
+        shapeless = True
         if axes is None:
             s = list(a.shape)
         else:
             s = take(a.shape, axes)
     else:
-        shapeless = 0
+        shapeless = False
     s = list(s)
     if axes is None:
+        if not shapeless and np.__version__ >= "2.0":
+            msg = (
+                "`axes` should not be `None` if `s` is not `None` "
+                "(Deprecated in NumPy 2.0). In a future version of NumPy, "
+                "this will raise an error and `s[i]` will correspond to "
+                "the size along the transformed axis specified by "
+                "`axes[i]`. To retain current behaviour, pass a sequence "
+                "[0, ..., k-1] to `axes` for an array of dimension k."
+            )
+            warnings.warn(msg, DeprecationWarning, stacklevel=3)
         axes = list(range(-len(s), 0))
     if len(s) != len(axes):
         raise ValueError("Shape and axes have different lengths.")
     if invreal and shapeless:
         s[-1] = (a.shape[axes[-1]] - 1) * 2
+    if None in s and np.__version__ >= "2.0":
+        msg = (
+            "Passing an array containing `None` values to `s` is "
+            "deprecated in NumPy 2.0 and will raise an error in "
+            "a future version of NumPy. To use the default behaviour "
+            "of the corresponding 1-D transform, pass the value matching "
+            "the default for its `n` parameter. To use the default "
+            "behaviour for every axis, the `s` argument can be omitted."
+        )
+        warnings.warn(msg, DeprecationWarning, stacklevel=3)
+    # use the whole input array along axis `i` if `s[i] == -1 or None`
+    s = [a.shape[_a] if _s in [-1, None] else _s for _s, _a in zip(s, axes)]
+
     return s, axes
 
 
@@ -806,15 +805,10 @@ def fftn(a, s=None, axes=None, norm=None):
     >>> plt.show()
 
     """
-    _check_norm(norm)
-    x = _float_utils.__downcast_float128_array(a)
 
-    if norm in (None, "backward"):
-        fsc = 1.0
-    elif norm == "forward":
-        fsc = frwd_sc_nd(s, x.shape)
-    else:
-        fsc = sqrt(frwd_sc_nd(s, x.shape))
+    x = _float_utils.__downcast_float128_array(a)
+    s, axes = _cook_nd_args(x, s, axes)
+    fsc = _compute_fwd_scale(norm, s, x.shape)
 
     return trycall(mkl_fft.fftn, (x,), {"s": s, "axes": axes, "fwd_scale": fsc})
 
@@ -918,15 +912,10 @@ def ifftn(a, s=None, axes=None, norm=None):
     >>> plt.show()
 
     """
-    _check_norm(norm)
-    x = _float_utils.__downcast_float128_array(a)
 
-    if norm in (None, "backward"):
-        fsc = 1.0
-    elif norm == "forward":
-        fsc = frwd_sc_nd(s, x.shape)
-    else:
-        fsc = sqrt(frwd_sc_nd(s, x.shape))
+    x = _float_utils.__downcast_float128_array(a)
+    s, axes = _cook_nd_args(x, s, axes)
+    fsc = _compute_fwd_scale(norm, s, x.shape)
 
     return trycall(
         mkl_fft.ifftn, (x,), {"s": s, "axes": axes, "fwd_scale": fsc}
@@ -1025,9 +1014,8 @@ def fft2(a, s=None, axes=(-2, -1), norm=None):
               0.0 +0.j        ,   0.0 +0.j        ]])
 
     """
-    _check_norm(norm)
-    x = _float_utils.__downcast_float128_array(a)
-    return fftn(x, s=s, axes=axes, norm=norm)
+
+    return fftn(a, s=s, axes=axes, norm=norm)
 
 
 def ifft2(a, s=None, axes=(-2, -1), norm=None):
@@ -1119,9 +1107,8 @@ def ifft2(a, s=None, axes=(-2, -1), norm=None):
            [ 0.+0.j,  1.+0.j,  0.+0.j,  0.+0.j]])
 
     """
-    _check_norm(norm)
-    x = _float_utils.__downcast_float128_array(a)
-    return ifftn(x, s=s, axes=axes, norm=norm)
+
+    return ifftn(a, s=s, axes=axes, norm=norm)
 
 
 def rfftn(a, s=None, axes=None, norm=None):
@@ -1213,19 +1200,10 @@ def rfftn(a, s=None, axes=None, norm=None):
             [ 0.+0.j,  0.+0.j]]])
 
     """
-    _check_norm(norm)
-    x = _float_utils.__downcast_float128_array(a)
 
-    if norm in (None, "backward"):
-        fsc = 1.0
-    elif norm == "forward":
-        x = asanyarray(x)
-        s, axes = _cook_nd_args(x, s, axes)
-        fsc = frwd_sc_nd(s, x.shape)
-    else:
-        x = asanyarray(x)
-        s, axes = _cook_nd_args(x, s, axes)
-        fsc = sqrt(frwd_sc_nd(s, x.shape))
+    x = _float_utils.__downcast_float128_array(a)
+    s, axes = _cook_nd_args(x, s, axes)
+    fsc = _compute_fwd_scale(norm, s, x.shape)
 
     return trycall(
         mkl_fft.rfftn, (x,), {"s": s, "axes": axes, "fwd_scale": fsc}
@@ -1271,9 +1249,8 @@ def rfft2(a, s=None, axes=(-2, -1), norm=None):
     For more details see `rfftn`.
 
     """
-    _check_norm(norm)
-    x = _float_utils.__downcast_float128_array(a)
-    return rfftn(x, s, axes, norm)
+
+    return rfftn(a, s, axes, norm)
 
 
 def irfftn(a, s=None, axes=None, norm=None):
@@ -1367,19 +1344,10 @@ def irfftn(a, s=None, axes=None, norm=None):
             [ 1.,  1.]]])
 
     """
-    _check_norm(norm)
-    x = _float_utils.__downcast_float128_array(a)
 
-    if norm in (None, "backward"):
-        fsc = 1.0
-    elif norm == "forward":
-        x = asanyarray(x)
-        s, axes = _cook_nd_args(x, s, axes, invreal=1)
-        fsc = frwd_sc_nd(s, x.shape)
-    else:
-        x = asanyarray(x)
-        s, axes = _cook_nd_args(x, s, axes, invreal=1)
-        fsc = sqrt(frwd_sc_nd(s, x.shape))
+    x = _float_utils.__downcast_float128_array(a)
+    s, axes = _cook_nd_args(x, s, axes, invreal=True)
+    fsc = _compute_fwd_scale(norm, s, x.shape)
 
     return trycall(
         mkl_fft.irfftn, (x,), {"s": s, "axes": axes, "fwd_scale": fsc}
@@ -1425,6 +1393,5 @@ def irfft2(a, s=None, axes=(-2, -1), norm=None):
     For more details see `irfftn`.
 
     """
-    _check_norm(norm)
-    x = _float_utils.__downcast_float128_array(a)
-    return irfftn(x, s, axes, norm)
+
+    return irfftn(a, s, axes, norm)

mkl_fft/_pydfti.pyx

@@ -1100,7 +1100,7 @@ def _fftnd_impl(x, s=None, axes=None, overwrite_x=False, direction=+1, double fs
                 _direct_fftnd,
                 {'overwrite_x': overwrite_x, 'direction': direction, 'fsc': fsc},
                 res
-                )
+            )
         else:
             sc = <object> fsc
             return _iter_fftnd(x, s=s, axes=axes,
@@ -1200,7 +1200,7 @@ def rfftn(x, s=None, axes=None, fwd_scale=1.0):
             a = _fix_dimensions(a, tuple(ss), axes)
         if len(set(axes)) == len(axes) and len(axes) == a.ndim and len(axes) > 2:
             ss, aa = _remove_axis(s, axes, -1)
-            ind = [slice(None,None,1),] * len(s)
+            ind = [slice(None, None, 1),] * len(s)
             for ii in range(a.shape[la]):
                 ind[la] = ii
                 tind = tuple(ind)