BUG: correct irfft with n=1 on larger input

mhvk · mhvk · commit faea9fbcaff1 · 2024-01-23T09:46:36.000-05:00
diff --git a/numpy/fft/_pocketfft_umath.c b/numpy/fft/_pocketfft_umath.c
@@ -33,6 +33,9 @@ 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 (ncopy <= 0) {
+        return;  /* safety, and indeed required for npts=1 in irfft */
+    }
     if (step_in == elsize && step_out == elsize) {
         memcpy(out, in, ncopy*elsize);
     }
@@ -235,11 +238,11 @@ irfft_loop(char **args, npy_intp const *dimensions, npy_intp const *steps, void
         double *op_double = (double *)op_or_buff;
         /* Copy complex input to buffer in FFTpack order */
         op_double[0] = ((double *)ip)[0];
-        if (nin > 1) {
+        if (npts > 1) {
             copy_data(ip + step_in, step_in, nin - 2,
                       (char *)&op_double[1], sizeof(npy_cdouble), npts / 2 - 1,
                       sizeof(npy_cdouble));
-            npy_intp ncopied = (npts / 2 - 1  < nin - 2 ? npts / 2 - 1 : nin - 2);
+            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 */
diff --git a/numpy/fft/tests/test_pocketfft.py b/numpy/fft/tests/test_pocketfft.py
@@ -33,6 +33,20 @@ 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_reversed(self):
+        maxlen = 16
+        x = random(maxlen) + 1j*random(maxlen)
+        for i in range(1, maxlen):
+            assert_allclose(np.fft.fft(np.fft.ifft(x, n=i)),
+                            x[0:i], atol=1e-12)
+            y = x.copy()
+            n = i // 2 + 1
+            y.imag[0] = 0
+            if i % 2 == 0:
+                y.imag[n-1] = 0
+            assert_allclose(np.fft.rfft(np.fft.irfft(y, n=i)),
+                            y[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 +506,10 @@ 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)
