more tests and fixes

pocketfft.cc

@@ -2282,7 +2282,7 @@ py::array irfftn(const py::array &in, vector<size_t> axes, size_t lastsize,
     inter = in;
 
   size_t axis = axes.back();
-  if (lastsize==0) lastsize=1+(inter.shape(axis)/2);
+  if (lastsize==0) lastsize=2*inter.shape(axis)-1;
   if (int64_t(lastsize/2) + 1 != inter.shape(axis))
     throw runtime_error("bad lastsize");
   vector<size_t> dims_in(inter.ndim()), dims_out(inter.ndim());
@@ -2365,6 +2365,8 @@ a : numpy.ndarray (np.complex64 or np.complex128)
 axes : list of integers
     The axes on which the FFT is carried out.
     If not set, all axes will be transformed.
+fct : float
+    normalization factor
 
 Returns
 -------
@@ -2382,6 +2384,8 @@ a : numpy.ndarray (np.complex64 or np.complex128)
 axes : list of integers
     The axes on which the FFT is carried out.
     If not set, all axes will be transformed.
+fct : float
+    normalization factor
 
 Returns
 -------
@@ -2389,6 +2393,50 @@ np.ndarray (same shape and data type as a)
     the transformed data
 )DELIM";
 
+const char *rfftn_DS = R"DELIM(
+Performs a forward real-valued FFT.
+
+Parameters
+----------
+a : numpy.ndarray (np.float32 or np.float64)
+    The input data
+axes : list of integers
+    The axes on which the FFT is carried out.
+    If not set, all axes will be transformed in ascending order.
+fct : float
+    normalization factor
+
+Returns
+-------
+np.ndarray (np.complex64 or np.complex128)
+    the transformed data. The shape is identical to that of the input array,
+    except for the axis that was transformed last. If the length of that axis
+    was n on input, it is n//2+1 on output.
+)DELIM";
+
+const char *irfftn_DS = R"DELIM(
+Performs a backward real-valued FFT.
+
+Parameters
+----------
+a : numpy.ndarray (np.complex64 or np.complex128)
+    The input data
+axes : list of integers
+    The axes on which the FFT is carried out.
+    If not set, all axes will be transformed in ascending order.
+lastsize : the output size of the last axis to be transformed.
+    If the corresponding input axis has size n, this can be 2*n-2 or 2*n-1.
+fct : float
+    normalization factor
+
+Returns
+-------
+np.ndarray (np.float32 or np.float64)
+    the transformed data. The shape is identical to that of the input array,
+    except for the axis that was transformed last, which has now lastsize
+    entries.
+)DELIM";
+
 const char *hartley_DS = R"DELIM(
 Performs a Hartley FFT.
 For every requested axis, a 1D forward Fourier transform is carried out,
@@ -2402,6 +2450,8 @@ a : numpy.ndarray (np.float32 or np.float64)
 axes : list of integers
     The axes on which the transform is carried out.
     If not set, all axes will be transformed.
+fct : float
+    normalization factor
 
 Returns
 -------
@@ -2417,7 +2467,7 @@ PYBIND11_MODULE(pypocketfft, m)
 
   m.def("fftn",&fftn, fftn_DS, "a"_a, "axes"_a=vector<size_t>(1,10000), "fct"_a=1.);
   m.def("ifftn",&ifftn, ifftn_DS, "a"_a, "axes"_a=vector<size_t>(1,10000), "fct"_a=1.);
-  m.def("rfftn",&rfftn, "a"_a, "axes"_a=vector<size_t>(1,10000), "fct"_a=1.);
-  m.def("irfftn",&irfftn, "a"_a, "axes"_a=vector<size_t>(1,10000), "lastsize"_a=0, "fct"_a=1.);
+  m.def("rfftn",&rfftn, rfftn_DS, "a"_a, "axes"_a=vector<size_t>(1,10000), "fct"_a=1.);
+  m.def("irfftn",&irfftn, irfftn_DS, "a"_a, "axes"_a=vector<size_t>(1,10000), "lastsize"_a=0, "fct"_a=1.);
   m.def("hartley",&hartley, hartley_DS, "a"_a, "axes"_a=vector<size_t>(1,10000), "fct"_a=1.);
   }

test.py

@@ -63,6 +63,14 @@ def test_identity_r(shp):
         assert_allclose(pypocketfft.irfftn(pypocketfft.rfftn(a,(ax,)),(ax,),lastsize=n,fct=1./n), a, atol=2e-15*vmax, rtol=0)
         assert_allclose(pypocketfft.irfftn(pypocketfft.rfftn(b,(ax,)),(ax,),lastsize=n,fct=1./n), b, atol=6e-7*vmax, rtol=0)
 
+@pmp("shp", shapes)
+def test_identity_r2(shp):
+    a=np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
+    a=pypocketfft.rfftn(pypocketfft.irfftn(a))
+    fct = 1./pypocketfft.irfftn(a).size
+    vmax = np.max(np.abs(a))
+    assert_allclose(pypocketfft.rfftn(pypocketfft.irfftn(a),fct=fct), a, atol=2e-15*vmax, rtol=0)
+
 @pmp("shp", shapes3D)
 def xtest_hartley(shp):
     a=np.random.rand(*shp)-0.5