API: adjust nD fft s param to array API

Author: lucascolley

numpy/fft/_pocketfft.py

@@ -681,20 +681,24 @@ def ihfft(a, n=None, axis=-1, norm=None):
 
 def _cook_nd_args(a, s=None, axes=None, invreal=0):
     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:
+            msg = "`axes` must not be `None` if `s` is not `None`."
+            raise ValueError(msg)
         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
+    s = [a.shape[_a] if _s == -1 else _s for _s, _a in zip(s, axes)]
     return s, axes
 
 
@@ -730,9 +734,11 @@ def fftn(a, s=None, axes=None, norm=None):
         (``s[0]`` refers to axis 0, ``s[1]`` to axis 1, etc.).
         This corresponds to ``n`` for ``fft(x, n)``.
         Along any axis, if the given shape is smaller than that of the input,
-        the input is cropped.  If it is larger, the input is padded with zeros.
-        if `s` is not given, the shape of the input along the axes specified
+        the input is cropped. If it is larger, the input is padded with zeros.
+        If it is ``-1``, the whole input is used (no padding/trimming).
+        If `s` is not given, the shape of the input along the axes specified
         by `axes` is used.
+        If `s` is not ``None``, `axes` must not be ``None`` either.
     axes : sequence of ints, optional
         Axes over which to compute the FFT.  If not given, the last ``len(s)``
         axes are used, or all axes if `s` is also not specified.
@@ -842,9 +848,11 @@ def ifftn(a, s=None, axes=None, norm=None):
         (``s[0]`` refers to axis 0, ``s[1]`` to axis 1, etc.).
         This corresponds to ``n`` for ``ifft(x, n)``.
         Along any axis, if the given shape is smaller than that of the input,
-        the input is cropped.  If it is larger, the input is padded with zeros.
-        if `s` is not given, the shape of the input along the axes specified
-        by `axes` is used.  See notes for issue on `ifft` zero padding.
+        the input is cropped. If it is larger, the input is padded with zeros.
+        If it is ``-1``, the whole input is used (no padding/trimming).
+        If `s` is not given, the shape of the input along the axes specified
+        by `axes` is used. See notes for issue on `ifft` zero padding.
+        If `s` is not ``None``, `axes` must not be ``None`` either.
     axes : sequence of ints, optional
         Axes over which to compute the IFFT.  If not given, the last ``len(s)``
         axes are used, or all axes if `s` is also not specified.
@@ -937,8 +945,9 @@ def fft2(a, s=None, axes=(-2, -1), norm=None):
         (``s[0]`` refers to axis 0, ``s[1]`` to axis 1, etc.).
         This corresponds to ``n`` for ``fft(x, n)``.
         Along each axis, if the given shape is smaller than that of the input,
-        the input is cropped.  If it is larger, the input is padded with zeros.
-        if `s` is not given, the shape of the input along the axes specified
+        the input is cropped. If it is larger, the input is padded with zeros.
+        If it is ``-1``, the whole input is used (no padding/trimming).
+        If `s` is not given, the shape of the input along the axes specified
         by `axes` is used.
     axes : sequence of ints, optional
         Axes over which to compute the FFT.  If not given, the last two
@@ -1040,8 +1049,9 @@ def ifft2(a, s=None, axes=(-2, -1), norm=None):
         Shape (length of each axis) of the output (``s[0]`` refers to axis 0,
         ``s[1]`` to axis 1, etc.).  This corresponds to `n` for ``ifft(x, n)``.
         Along each axis, if the given shape is smaller than that of the input,
-        the input is cropped.  If it is larger, the input is padded with zeros.
-        if `s` is not given, the shape of the input along the axes specified
+        the input is cropped. If it is larger, the input is padded with zeros.
+        If it is ``-1``, the whole input is used (no padding/trimming).
+        If `s` is not given, the shape of the input along the axes specified
         by `axes` is used.  See notes for issue on `ifft` zero padding.
     axes : sequence of ints, optional
         Axes over which to compute the FFT.  If not given, the last two
@@ -1128,9 +1138,11 @@ def rfftn(a, s=None, axes=None, norm=None):
         The final element of `s` corresponds to `n` for ``rfft(x, n)``, while
         for the remaining axes, it corresponds to `n` for ``fft(x, n)``.
         Along any axis, if the given shape is smaller than that of the input,
-        the input is cropped.  If it is larger, the input is padded with zeros.
-        if `s` is not given, the shape of the input along the axes specified
+        the input is cropped. If it is larger, the input is padded with zeros.
+        If it is ``-1``, the whole input is used (no padding/trimming).
+        If `s` is not given, the shape of the input along the axes specified
         by `axes` is used.
+        If `s` is not ``None``, `axes` must not be ``None`` either.
     axes : sequence of ints, optional
         Axes over which to compute the FFT.  If not given, the last ``len(s)``
         axes are used, or all axes if `s` is also not specified.
@@ -1284,9 +1296,12 @@ def irfftn(a, s=None, axes=None, norm=None):
         where ``s[-1]//2+1`` points of the input are used.
         Along any axis, if the shape indicated by `s` is smaller than that of
         the input, the input is cropped.  If it is larger, the input is padded
-        with zeros. If `s` is not given, the shape of the input along the axes
+        with zeros.
+        If it is ``-1``, the whole input is used (no padding/trimming).
+        If `s` is not given, the shape of the input along the axes
         specified by axes is used. Except for the last axis which is taken to
         be ``2*(m-1)`` where ``m`` is the length of the input along that axis.
+        If `s` is not ``None``, `axes` must not be ``None`` either.
     axes : sequence of ints, optional
         Axes over which to compute the inverse FFT. If not given, the last
         `len(s)` axes are used, or all axes if `s` is also not specified.

numpy/fft/tests/test_pocketfft.py

@@ -199,6 +199,19 @@ def test_axes(self, op):
             op_tr = op(np.transpose(x, a))
             tr_op = np.transpose(op(x, axes=a), a)
             assert_allclose(op_tr, tr_op, atol=1e-6)
+    
+    @pytest.mark.parametrize("op", [np.fft.fftn, np.fft.ifftn])
+    def test_s_negative_1(self, op):
+        x = np.arange(100).reshape(10, 10)
+        # should use the whole input array along the first axis
+        assert op(x, s=(-1, 5), axes=(0, 1)).shape == (10, 5)
+
+    @pytest.mark.parametrize("op", [np.fft.fftn, np.fft.ifftn,
+                                    np.fft.rfftn, np.fft.irfftn])
+    def test_s_axes_none(self, op):
+        x = np.arange(100).reshape(10, 10)
+        with pytest.raises(ValueError):
+            op(x, s=(-1, 5))
 
     def test_all_1d_norm_preserving(self):
         # verify that round-trip transforms are norm-preserving