Merge pull request #69 from cako/master

Fixes NumPy interface bug which did not allow for list/ndarray `axes` or `s` parameters

Author: oleksandr-pavlyk

mkl_fft/_numpy_fft.py

@@ -76,8 +76,8 @@ def frwd_sc_1d(n, s):
 
 
 def frwd_sc_nd(s, axes, x_shape):
-    ss = s if s else x_shape
-    if axes:
+    ss = s if s is not None else x_shape
+    if axes is not None:
         nn = prod([ss[ai] for ai in axes])
     else:
         nn = prod(ss)
@@ -203,7 +203,7 @@ def fft(a, n=None, axis=-1, norm=None):
             mkl_fft.fft,
             (x,),
             {'n':n, 'axis': axis})
-    elif norm is "forward":
+    elif norm == "forward":
         output = trycall(
             mkl_fft.fft,
             (x,),

mkl_fft/_pydfti.pyx

@@ -932,7 +932,7 @@ def iter_complementary(x, axes, func, kwargs, result):
     nd = x.ndim
     r = list(range(nd))
     sl = [slice(None, None, None)] * nd
-    if not isinstance(axes, tuple):
+    if not np.iterable(axes):
         axes = (axes,)
     for ai in axes:
         r[ai] = None

mkl_fft/tests/test_fftnd.py

@@ -114,6 +114,18 @@ def test_matrix5(self):
                     rtol=r_tol, atol=a_tol
                 )
 
+    def test_matrix6(self):
+        """fftn with tuple, list and ndarray axes and s"""
+        for ar in [self.md, self.mz, self.mf, self.mc]:
+            d = ar.copy()
+            for norm in ["forward", "backward", "ortho"]:
+                for container in [tuple, list, np.array]:
+                    axes = container(range(d.ndim))
+                    s = container(d.shape)
+                    kwargs = dict(s=s, axes=axes, norm=norm)
+                    r_tol, a_tol = _get_rtol_atol(d)
+                    t = mkl_fft._numpy_fft.fftn(mkl_fft._numpy_fft.ifftn(d, **kwargs), **kwargs)
+                    assert_allclose(d, t, rtol=r_tol, atol=a_tol, err_msg = "failed test for dtype {}, max abs diff: {}".format(d.dtype, np.max(np.abs(d-t))))