added support for forward_scale parameter, and tests

Author: oleksandr-pavlyk

mkl_fft/_pydfti.pyx

@@ -151,12 +151,12 @@ cdef int _datacopied(cnp.ndarray arr, object orig):
     return 1 if (arr_obj.base is None) else 0
 
 
-def fft(x, n=None, axis=-1, overwrite_x=False):
-    return _fft1d_impl(x, n=n, axis=axis, overwrite_arg=overwrite_x, direction=+1)
+def fft(x, n=None, axis=-1, overwrite_x=False, forward_scale=1.0):
+    return _fft1d_impl(x, n=n, axis=axis, overwrite_arg=overwrite_x, direction=+1, fsc=forward_scale)
 
 
-def ifft(x, n=None, axis=-1, overwrite_x=False):
-    return _fft1d_impl(x, n=n, axis=axis, overwrite_arg=overwrite_x, direction=-1)
+def ifft(x, n=None, axis=-1, overwrite_x=False, forward_scale=1.0):
+    return _fft1d_impl(x, n=n, axis=axis, overwrite_arg=overwrite_x, direction=-1, fsc=forward_scale)
 
 
 cdef cnp.ndarray pad_array(cnp.ndarray x_arr, cnp.npy_intp n, int axis, int realQ):
@@ -403,14 +403,14 @@ def _fft1d_impl(x, n=None, axis=-1, overwrite_arg=False, direction=+1, double fs
         return f_arr
 
 
-def rfft(x, n=None, axis=-1, overwrite_x=False):
+def rfft(x, n=None, axis=-1, overwrite_x=False, forward_scale=1.0):
     """Packed real-valued harmonics of FFT of a real sequence x"""
-    return _rr_fft1d_impl2(x, n=n, axis=axis, overwrite_arg=overwrite_x)
+    return _rr_fft1d_impl2(x, n=n, axis=axis, overwrite_arg=overwrite_x, fsc=forward_scale)
 
 
-def irfft(x, n=None, axis=-1, overwrite_x=False):
+def irfft(x, n=None, axis=-1, overwrite_x=False, forward_scale=1.0):
     """Inverse FFT of a real sequence, takes packed real-valued harmonics of FFT"""
-    return _rr_ifft1d_impl2(x, n=n, axis=axis, overwrite_arg=overwrite_x)
+    return _rr_ifft1d_impl2(x, n=n, axis=axis, overwrite_arg=overwrite_x, fsc=forward_scale)
 
 
 cdef object _rc_to_rr(cnp.ndarray rc_arr, int n, int axis, int xnd, int x_type):
@@ -788,12 +788,12 @@ def _rc_ifft1d_impl(x, n=None, axis=-1, overwrite_arg=False, double fsc=1.0):
         return f_arr
 
 
-def rfft_numpy(x, n=None, axis=-1):
-    return _rc_fft1d_impl(x, n=n, axis=axis)
+def rfft_numpy(x, n=None, axis=-1, forward_scale=1.0):
+    return _rc_fft1d_impl(x, n=n, axis=axis, fsc=forward_scale)
 
 
-def irfft_numpy(x, n=None, axis=-1):
-    return _rc_ifft1d_impl(x, n=n, axis=axis)
+def irfft_numpy(x, n=None, axis=-1, forward_scale=1.0):
+    return _rc_ifft1d_impl(x, n=n, axis=axis, fsc=forward_scale)
 
 
 # ============================== ND ====================================== #
@@ -902,12 +902,12 @@ def _cook_nd_args(a, s=None, axes=None, invreal=0):
     return s, axes
 
 
-def _iter_fftnd(a, s=None, axes=None, function=fft, overwrite_arg=False):
+def _iter_fftnd(a, s=None, axes=None, function=fft, overwrite_arg=False, scale_function=lambda n: 1.0):
     a = np.asarray(a)
     s, axes = _init_nd_shape_and_axes(a, s, axes)
     ovwr = overwrite_arg
     for ii in reversed(range(len(axes))):
-        a = function(a, n = s[ii], axis = axes[ii], overwrite_x=ovwr)
+        a = function(a, n = s[ii], axis = axes[ii], overwrite_x=ovwr, forward_scale=scale_function(s[ii]))
         ovwr = True
     return a
 
@@ -1026,33 +1026,34 @@ def _fftnd_impl(x, shape=None, axes=None, overwrite_x=False, direction=+1, doubl
     if _direct:
         return _direct_fftnd(x, overwrite_arg=overwrite_x, direction=direction, fsc=fsc)
     else:
+        sc = (<object> fsc)**(1/x.ndim)
         return _iter_fftnd(x, s=shape, axes=axes,
-                           overwrite_arg=overwrite_x, fsc=fsc,
+                           overwrite_arg=overwrite_x, scale_function=lambda n: sc,
                            function=fft if direction == 1 else ifft)
 
 
-def fft2(x, shape=None, axes=(-2,-1), overwrite_x=False):
-    return _fftnd_impl(x, shape=shape, axes=axes, overwrite_x=overwrite_x, direction=+1)
+def fft2(x, shape=None, axes=(-2,-1), overwrite_x=False, forward_scale=1.0):
+    return _fftnd_impl(x, shape=shape, axes=axes, overwrite_x=overwrite_x, direction=+1, fsc=forward_scale)
 
 
-def ifft2(x, shape=None, axes=(-2,-1), overwrite_x=False):
-    return _fftnd_impl(x, shape=shape, axes=axes, overwrite_x=overwrite_x, direction=-1)
+def ifft2(x, shape=None, axes=(-2,-1), overwrite_x=False, forward_scale=1.0):
+    return _fftnd_impl(x, shape=shape, axes=axes, overwrite_x=overwrite_x, direction=-1, fsc=forward_scale)
 
 
-def fftn(x, shape=None, axes=None, overwrite_x=False):
-    return _fftnd_impl(x, shape=shape, axes=axes, overwrite_x=overwrite_x, direction=+1)
+def fftn(x, shape=None, axes=None, overwrite_x=False, forward_scale=1.0):
+    return _fftnd_impl(x, shape=shape, axes=axes, overwrite_x=overwrite_x, direction=+1, fsc=forward_scale)
 
 
-def ifftn(x, shape=None, axes=None, overwrite_x=False):
-    return _fftnd_impl(x, shape=shape, axes=axes, overwrite_x=overwrite_x, direction=-1)
+def ifftn(x, shape=None, axes=None, overwrite_x=False, forward_scale=1.0):
+    return _fftnd_impl(x, shape=shape, axes=axes, overwrite_x=overwrite_x, direction=-1, fsc=forward_scale)
 
 
-def rfft2_numpy(x, s=None, axes=(-2,-1)):
-    return rfftn_numpy(x, s=s, axes=axes)
+def rfft2_numpy(x, s=None, axes=(-2,-1), forward_scale=1.0):
+    return rfftn_numpy(x, s=s, axes=axes, fsc=forward_scale)
 
 
-def irfft2_numpy(x, s=None, axes=(-2,-1)):
-    return irfftn_numpy(x, s=s, axes=axes)
+def irfft2_numpy(x, s=None, axes=(-2,-1), forward_scale=1.0):
+    return irfftn_numpy(x, s=s, axes=axes, fsc=forward_scale)
 
 
 def _remove_axis(s, axes, axis_to_remove):
@@ -1107,7 +1108,7 @@ def _fix_dimensions(cnp.ndarray arr, object s, object axes):
     return np.pad(arr, tuple(pad_widths), 'constant')
 
 
-def rfftn_numpy(x, s=None, axes=None):
+def rfftn_numpy(x, s=None, axes=None, forward_scale=1.0):
     a = np.asarray(x)
     no_trim = (s is None) and (axes is None)
     s, axes = _cook_nd_args(a, s, axes)
@@ -1116,7 +1117,7 @@ def rfftn_numpy(x, s=None, axes=None):
     # unnecessary computations
     if not no_trim:
         a = _trim_array(a, s, axes)
-    a = rfft_numpy(a, n = s[-1], axis=la)
+    a = rfft_numpy(a, n = s[-1], axis=la, forward_scale=forward_scale)
     if len(s) > 1:
         if not no_trim:
             ss = list(s)
@@ -1140,7 +1141,7 @@ def rfftn_numpy(x, s=None, axes=None):
     return a
 
 
-def irfftn_numpy(x, s=None, axes=None):
+def irfftn_numpy(x, s=None, axes=None, forward_scale=1.0):
     a = np.asarray(x)
     no_trim = (s is None) and (axes is None)
     s, axes = _cook_nd_args(a, s, axes, invreal=True)
@@ -1169,5 +1170,5 @@ def irfftn_numpy(x, s=None, axes=None):
             for ii in range(len(axes)-1):
                 a = ifft(a, s[ii], axes[ii], overwrite_x=ovr_x)
                 ovr_x = True
-    a = irfft_numpy(a, n = s[-1], axis=la)
+    a = irfft_numpy(a, n = s[-1], axis=la, forward_scale=forward_scale)
     return a

mkl_fft/tests/test_fftnd.py

@@ -38,17 +38,17 @@
 
 reps_64 = (2**11)*np.finfo(np.float64).eps
 reps_32 = (2**11)*np.finfo(np.float32).eps
-atol_64 = (2**8)*np.finfo(np.float64).eps
-atol_32 = (2**8)*np.finfo(np.float32).eps
+atol_64 = (2**9)*np.finfo(np.float64).eps
+atol_32 = (2**9)*np.finfo(np.float32).eps
 
 def _get_rtol_atol(x):
     dt = x.dtype
-    if dt == np.double or dt == np.complex128:
+    if dt == np.float64 or dt == np.complex128:
         return reps_64, atol_64
-    elif dt == np.single or dt == np.complex64:
+    elif dt == np.float32 or dt == np.complex64:
         return reps_32, atol_32
     else:
-        assert (dt == np.double or dt == np.complex128 or dt == np.single or dt == np.complex64), "Unexpected dtype {}".format(dt)
+        assert (dt == np.float64 or dt == np.complex128 or dt == np.float32 or dt == np.complex64), "Unexpected dtype {}".format(dt)
         return reps_64, atol_64
 
 
@@ -128,3 +128,54 @@ def test_rfftn_numpy(self):
                 rfft_tr = mkl_fft.rfftn_numpy(np.transpose(x, a))
                 tr_rfft = np.transpose(mkl_fft.rfftn_numpy(x, axes=a), a)
                 assert_allclose(rfft_tr, tr_rfft, rtol=r_tol, atol=a_tol)
+
+class Test_Scales(TestCase):
+    def setUp(self):
+        pass
+
+    def test_scale_1d_vector(self):
+        X = np.ones(128, dtype='d')
+        f1 = mkl_fft.fft(X, forward_scale=0.25)
+        f2 = mkl_fft.fft(X)
+        r_tol, a_tol = _get_rtol_atol(X)
+        assert_allclose(4*f1, f2, rtol=r_tol, atol=a_tol)
+
+        X1 = mkl_fft.ifft(f1, forward_scale=0.25)
+        assert_allclose(X, X1, rtol=r_tol, atol=a_tol)
+
+        f3 = mkl_fft.rfft(X, forward_scale=0.5)
+        X2 = mkl_fft.irfft(f3, forward_scale=0.5)
+        assert_allclose(X, X2, rtol=r_tol, atol=a_tol)
+
+    def test_scale_1d_array(self):
+        X = np.ones((8, 4, 4,), dtype='d')
+        f1 = mkl_fft.fft(X, axis=1, forward_scale=0.25)
+        f2 = mkl_fft.fft(X, axis=1)
+        r_tol, a_tol = _get_rtol_atol(X)
+        assert_allclose(4*f1, f2, rtol=r_tol, atol=a_tol)
+
+        X1 = mkl_fft.ifft(f1, axis=1, forward_scale=0.25)
+        assert_allclose(X, X1, rtol=r_tol, atol=a_tol)
+
+        f3 = mkl_fft.rfft(X, axis=0, forward_scale=0.5)
+        X2 = mkl_fft.irfft(f3, axis=0, forward_scale=0.5)
+        assert_allclose(X, X2, rtol=r_tol, atol=a_tol)
+
+    def test_scale_nd(self):
+        X = np.empty((2, 4, 8, 16), dtype='d')
+        X.flat[:] = np.cbrt(np.arange(0, X.size, dtype=X.dtype))
+        f = mkl_fft.fftn(X)
+        f_scale = mkl_fft.fftn(X, forward_scale=0.2)
+        
+        r_tol, a_tol = _get_rtol_atol(X)
+        assert_allclose(f, 5*f_scale, rtol=r_tol, atol=a_tol)
+
+    def test_scale_nd_axes(self):
+        X = np.empty((4, 2, 16, 8), dtype='d')
+        X.flat[:] = np.cbrt(np.arange(X.size, dtype=X.dtype))
+        f = mkl_fft.fftn(X, axes=(0, 1, 2, 3))
+        f_scale = mkl_fft.fftn(X, axes=(0, 1, 2, 3), forward_scale=0.2)
+        
+        r_tol, a_tol = _get_rtol_atol(X)
+        assert_allclose(f, 5*f_scale, rtol=r_tol, atol=a_tol)
+