better docstrings; more tests

pypocketfft.cc

@@ -331,8 +331,9 @@ Performs a forward complex FFT.
 
 Parameters
 ----------
-a : numpy.ndarray (np.complex64 or np.complex128)
-    The input data
+a : numpy.ndarray (any complex or real type)
+    The input data. If its type is real, a more efficient real-to-complex
+    transform will be used.
 axes : list of integers
     The axes along which the FFT is carried out.
     If not set, all axes will be transformed.
@@ -345,13 +346,15 @@ inorm : int
 inplace : bool
     if False, returns the result in a new array and leaves the input unchanged.
     if True, stores the result in the input array and returns a handle to it.
+    NOTE: if `a` is real-valued and `inplace` is `True`, an exception will be
+    raised!
 nthreads : int
     Number of threads to use. If 0, use the system default (typically governed
     by the `OMP_NUM_THREADS` environment variable).
 
 Returns
 -------
-np.ndarray (same shape and data type as a)
+np.ndarray (same shape as `a`, complex type with same accuracy as `a`)
     The transformed data.
 )""";
 
@@ -359,7 +362,7 @@ const char *ifftn_DS = R"""(Performs a backward complex FFT.
 
 Parameters
 ----------
-a : numpy.ndarray (np.complex64 or np.complex128)
+a : numpy.ndarray (any complex type)
     The input data
 axes : list of integers
     The axes along which the FFT is carried out.
@@ -387,7 +390,7 @@ const char *rfftn_DS = R"""(Performs a forward real-valued FFT.
 
 Parameters
 ----------
-a : numpy.ndarray (np.float32 or np.float64)
+a : numpy.ndarray (any real type)
     The input data
 axes : list of integers
     The axes along which the FFT is carried out.
@@ -404,7 +407,7 @@ nthreads : int
 
 Returns
 -------
-np.ndarray (np.complex64 or np.complex128)
+np.ndarray (complex type with same accuracy as `a`)
     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.
@@ -414,7 +417,7 @@ const char *rfft_scipy_DS = R"""(Performs a forward real-valued FFT.
 
 Parameters
 ----------
-a : numpy.ndarray (np.float32 or np.float64)
+a : numpy.ndarray (any real type)
     The input data
 axis : int
     The axis along which the FFT is carried out.
@@ -433,7 +436,7 @@ nthreads : int
 
 Returns
 -------
-np.ndarray (np.float32 or np.float64)
+np.ndarray (same shape and data type as `a`)
     The transformed data. The shape is identical to that of the input array.
     Along the transformed axis, values are arranged in
     FFTPACK half-complex order, i.e. `a[0].re, a[1].re, a[1].im, a[2].re ...`.
@@ -443,7 +446,7 @@ const char *irfftn_DS = R"""(Performs a backward real-valued FFT.
 
 Parameters
 ----------
-a : numpy.ndarray (np.complex64 or np.complex128)
+a : numpy.ndarray (any complex type)
     The input data
 axes : list of integers
     The axes along which the FFT is carried out.
@@ -462,7 +465,7 @@ nthreads : int
 
 Returns
 -------
-np.ndarray (np.float32 or np.float64)
+np.ndarray (real type with same accuracy as `a`)
     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.
@@ -472,7 +475,7 @@ const char *irfft_scipy_DS = R"""(Performs a backward real-valued FFT.
 
 Parameters
 ----------
-a : numpy.ndarray (np.float32 or np.float64)
+a : numpy.ndarray (any real type)
     The input data. Along the transformed axis, values are expected in
     FFTPACK half-complex order, i.e. `a[0].re, a[1].re, a[1].im, a[2].re ...`.
 axis : int
@@ -492,7 +495,7 @@ nthreads : int
 
 Returns
 -------
-np.ndarray (np.float32 or np.float64)
+np.ndarray (same shape and data type as `a`)
     The transformed data. The shape is identical to that of the input array.
 )""";
 
@@ -523,7 +526,7 @@ nthreads : int
 
 Returns
 -------
-np.ndarray (same shape and data type as a)
+np.ndarray (same shape and data type as `a`)
     The transformed data
 )""";
 

test.py

@@ -23,11 +23,13 @@ def test1D(len, inorm):
     c=a.astype(np.complex256)
     assert_(_l2error(a, pypocketfft.ifftn(pypocketfft.fftn(c,inorm=inorm), inorm=2-inorm))<1e-18)
     assert_(_l2error(a, pypocketfft.ifftn(pypocketfft.fftn(a,inorm=inorm), inorm=2-inorm))<1.5e-15)
+    assert_(_l2error(a.real, pypocketfft.ifftn(pypocketfft.fftn(a.real,inorm=inorm), inorm=2-inorm))<1.5e-15)
     assert_(_l2error(a.real, pypocketfft.irfftn(pypocketfft.rfftn(a.real,inorm=inorm), inorm=2-inorm,lastsize=len))<1.5e-15)
     tmp=a.copy()
     assert_ (pypocketfft.ifftn(pypocketfft.fftn(tmp, inplace=True, inorm=inorm), inplace=True, inorm=2-inorm) is tmp)
     assert_(_l2error(tmp, a)<1.5e-15)
     assert_(_l2error(b, pypocketfft.ifftn(pypocketfft.fftn(b, inorm=inorm), inorm=2-inorm))<6e-7)
+    assert_(_l2error(b.real, pypocketfft.ifftn(pypocketfft.fftn(b.real,inorm=inorm), inorm=2-inorm))<6e-7)
     assert_(_l2error(b.real, pypocketfft.irfftn(pypocketfft.rfftn(b.real, inorm=inorm), lastsize=len, inorm=2-inorm))<6e-7)
     tmp=b.copy()
     assert_ (pypocketfft.ifftn(pypocketfft.fftn(tmp, inplace=True, inorm=inorm), inplace=True, inorm=2-inorm) is tmp)