numpy · mhvk · Jan 25, 2024 · Jan 23, 2024 · Jan 23, 2024
diff --git a/doc/release/upcoming_changes/25536.new_feature.rst b/doc/release/upcoming_changes/25536.new_feature.rst
@@ -0,0 +1,4 @@
+``out`` for `numpy.fft`
+-----------------------
+The various FFT routines in `numpy.fft` have gained an ``out``
+argument that can be used for in-place calculations.
diff --git a/numpy/fft/_pocketfft_umath.c b/numpy/fft/_pocketfft_umath.c
@@ -33,22 +33,27 @@ copy_data(char* in, npy_intp step_in, npy_intp nin,
           char* out, npy_intp step_out, npy_intp nout, npy_intp elsize)
 {
     npy_intp ncopy = nin <= nout? nin : nout;
-    if (step_in == elsize && step_out == elsize) {
-        memcpy(out, in, ncopy*elsize);
+    if (ncopy > 0) {
+        if (step_in == elsize && step_out == elsize) {
+            memcpy(out, in, ncopy*elsize);
+        }
+        else {
+            char *ip = in, *op = out;
+            for (npy_intp i = 0; i < ncopy; i++, ip += step_in, op += step_out) {
+                memcpy(op, ip, elsize);
+            }
+        }
     }
     else {
-        char *ip = in, *op = out;
-        for (npy_intp i = 0; i < ncopy; i++, ip += step_in, op += step_out) {
-            memcpy(op, ip, elsize);
-        }
+        ncopy = 0;  /* can be negative, from irfft */
     }
-    if (nin < nout) {
-        char *op = out + nin*elsize;
+    if (nout > ncopy) {
+        char *op = out + ncopy*elsize;
         if (step_out == elsize) {
-            memset(op, 0, (nout-nin)*elsize);
+            memset(op, 0, (nout-ncopy)*elsize);
         }
         else {
-            for (npy_intp i = nin; i < nout; i++, op += step_out) {
+            for (npy_intp i = ncopy; i < nout; i++, op += step_out) {
                 memset(op, 0, elsize);
             }
         }
@@ -95,7 +100,11 @@ fft_loop(char **args, npy_intp const *dimensions, npy_intp const *steps, void *f
     for (npy_intp i = 0; i < n_outer; i++, ip += si, fp += sf, op += so) {
         double fct = *(double *)fp;
         char *op_or_buff = buff == NULL ? op : buff;
-        if (ip != op_or_buff) {  /* no copy needed if in-place already */
+        /*
+         * pocketfft works in-place, so we need to copy the data
+         * (except if we want to be in-place)
+         */
+        if (ip != op_or_buff) {
             copy_data(ip, step_in, nin,
                       op_or_buff, sizeof(npy_cdouble), npts, sizeof(npy_cdouble));
         }
@@ -154,7 +163,17 @@ rfft_impl(char **args, npy_intp const *dimensions, npy_intp const *steps, void *
         double fct = *(double *)fp;
         char *op_or_buff = buff == NULL ? op : buff;
         double *op_double = (double *)op_or_buff;
-        /* Copy dat to buffer in starting at position 1, as expected for FFTpack order */
+        /*
+         * Pocketfft works in-place and for real transforms the frequency data
+         * thus needs to be compressed, using that there will be no imaginary
+         * component for the zero-frequency item (which is the sum of all
+         * inputs and thus has to be real), nor one for the Nyquist frequency
+         * for even number of points. Pocketfft uses FFTpack order,
+         * R0,R1,I1,...Rn-1,In-1,Rn[,In] (last for npts odd only). To make
+         * unpacking easy, we place the real data offset by one in the buffer,
+         * so that we just have to move R0 and create I0=0. Note that
+         * copy_data will zero the In component for even number of points.
+         */
         copy_data(ip, step_in, nin,
                   (char *)&op_double[1], sizeof(npy_double), nout*2 - 1, sizeof(npy_double));
         if ((no_mem = rfft_forward(plan, &op_double[1], fct)) != 0) {
@@ -233,17 +252,29 @@ irfft_loop(char **args, npy_intp const *dimensions, npy_intp const *steps, void
         double fct = *(double *)fp;
         char *op_or_buff = buff == NULL ? op : buff;
         double *op_double = (double *)op_or_buff;
-        /* Copy complex input to buffer in FFTpack order */
-        op_double[0] = ((double *)ip)[0];
-        if (nin > 1) {
-            copy_data(ip + step_in, step_in, nin - 2,
-                      (char *)&op_double[1], sizeof(npy_cdouble), npts / 2 - 1,
+        /*
+         * Pocket_fft works in-place and for inverse real transforms the
+         * frequency data thus needs to be compressed, removing the imaginary
+         * component of the zero-frequency item (which is the sum of all
+         * inputs and thus has to be real), as well as the imaginary component
+         * of the Nyquist frequency for even number of points. We thus copy
+         * the data to the buffer in the following order (also used by
+         * FFTpack): R0,R1,I1,...Rn-1,In-1,Rn[,In] (last for npts odd only).
+         */
+        op_double[0] = ((double *)ip)[0];  /* copy R0 */
+        if (npts > 1) {
+            /*
+             * Copy R1,I1... up to Rn-1,In-1 if possible, stopping earlier
+             * if not all the input points are needed or if the input is short
+             * (in the latter case, zeroing after).
+             */
+            copy_data(ip + step_in, step_in, nin - 1,
+                      (char *)&op_double[1], sizeof(npy_cdouble), (npts - 1) / 2,
                       sizeof(npy_cdouble));
-            npy_intp ncopied = (npts / 2 - 1  < nin - 2 ? npts / 2 - 1 : nin - 2);
-            double *last = (double *)(ip + (ncopied + 1) * step_in);
-            op_double[ncopied * 2 + 1] = last[0];
-            if (npts % 2 == 1) {  /* For odd n, we still imag real of the last point */
-                op_double[ncopied * 2 + 2] = last[1];
+            /* For even npts, we still need to set Rn. */
+            if (npts % 2 == 0) {
+                op_double[npts - 1] = (npts / 2 >= nin) ? 0. :
+                    ((double *)(ip + (npts / 2) * step_in))[0];
             }
         }
         if ((no_mem = rfft_backward(plan, op_double, fct)) != 0) {

diff --git a/numpy/fft/tests/test_pocketfft.py b/numpy/fft/tests/test_pocketfft.py
@@ -33,6 +33,40 @@ def test_identity(self):
             assert_allclose(np.fft.irfft(np.fft.rfft(xr[0:i]), i),
                             xr[0:i], atol=1e-12)
 
+    def test_identity_long_short(self):
+        # Test with explicitly given number of points, both for n
+        # smaller and for n larger than the input size.
+        maxlen = 16
+        x = random(maxlen) + 1j*random(maxlen)
+        xx = np.concatenate([x, np.zeros_like(x)])
+        xr = random(maxlen)
+        xxr = np.concatenate([xr, np.zeros_like(xr)])
+        for i in range(1, maxlen*2):
+            assert_allclose(np.fft.ifft(np.fft.fft(x, n=i), n=i),
+                            xx[0:i], atol=1e-12)
+            assert_allclose(np.fft.irfft(np.fft.rfft(xr, n=i), n=i),
+                            xxr[0:i], atol=1e-12)
+
+    def test_identity_long_short_reversed(self):
+        # Also test explicitly given number of points in reversed order.
+        maxlen = 16
+        x = random(maxlen) + 1j*random(maxlen)
+        xx = np.concatenate([x, np.zeros_like(x)])
+        for i in range(1, maxlen*2):
+            assert_allclose(np.fft.fft(np.fft.ifft(x, n=i), n=i),
+                            xx[0:i], atol=1e-12)
+            # For irfft, we can neither recover the imaginary part of
+            # the first element, nor the imaginary part of the last
+            # element if npts is even.  So, set to 0 for the comparison.
+            y = x.copy()
+            n = i // 2 + 1
+            y.imag[0] = 0
+            if i % 2 == 0:
+                y.imag[n-1:] = 0
+            yy = np.concatenate([y, np.zeros_like(y)])
+            assert_allclose(np.fft.rfft(np.fft.irfft(x, n=i), n=i),
+                            yy[0:n], atol=1e-12)
+
     def test_fft(self):
         x = random(30) + 1j*random(30)
         assert_allclose(fft1(x), np.fft.fft(x), atol=1e-6)
@@ -492,3 +526,21 @@ def test_rfft(self):
     def test_irfft(self):
         a = np.ones(self.input_shape) * 1+0j
         self._test_mtsame(np.fft.irfft, a)
+
+
+def test_irfft_with_n_1_regression():
+    # Regression test for gh-25661
+    x = np.arange(10)
+    np.fft.irfft(x, n=1)
+    np.fft.hfft(x, n=1)
+    np.fft.irfft(np.array([0], complex), n=10)
+
+
+def test_irfft_with_n_large_regression():
+    # Regression test for gh-25679
+    x = np.arange(5) * (1 + 1j)
+    result = np.fft.hfft(x, n=10)
+    expected = np.array([20., 9.91628173, -11.8819096, 7.1048486,
+                         -6.62459848, 4., -3.37540152, -0.16057669,
+                         1.8819096, -20.86055364])
+    assert_allclose(result, expected)