more elegant implementation of fftfreqn

fresnel/misc.py

@@ -2,28 +2,42 @@ import numpy as np
 from scipy.fft import fftfreq
 
 
-def fftfreqn(N, d=1.0):
-    ndim = len(N)
+def fftfreqn(n, dx=1):
+    """
+    Return the nd Discrete Fourier Transform sample frequencies.
 
-    d *= np.ones(ndim)
+    Parameters
+    ----------
+    n : int or tuple of ints
+        Window lengths.
+    dx: float or tuple floats, optional (Default = 1.0)
+        Sample spacings.
 
-    # index trick: n'th line is a list of ones with a -1 on n'th position
-    shapes = np.ones((ndim, ndim), dtype="int") - 2 * np.eye(ndim, dtype="int")
+    Returns
+    -------
+    xi1, xi2, ..., xin : ndarray
 
-    xi = [np.reshape(fftfreq(N[dim], d[dim]), tuple(shapes[dim, :])) for dim in range(ndim)]
+    See Also
+    --------
+    scipy.fft.fftfreq
+    """
+    n = np.asarray(n)
+    ndim = n.size
 
-    return xi
+    dx = np.ones(ndim, dtype=float) * np.asarray(dx)
 
+    xi = [fftfreq(n_i, dx_i) for n_i, dx_i in zip(n, dx)]
 
-def gridn(N):
-    ndim = len(N)
+    return np.meshgrid(*xi, sparse=True)
 
-    # index trick: n'th line is a list of ones with a -1 on n'th position
-    shapes = np.ones((ndim, ndim), dtype="int") - 2 * np.eye(ndim, dtype="int")
 
-    x = [np.reshape(np.arange(-N[dim] + 1, N[dim]), tuple(shapes[dim, :])) for dim in range(ndim)]
+def gridn(n):
+    n = np.asarray(n)
+    ndim = n.size
 
-    return x
+    x = [np.arange(-n[dim] + 1, n[dim]) for dim in range(ndim)]
+
+    return np.meshgrid(*x, sparse=True)
 
 
 def crop(a, crop_width):