Commit c55aa9af authored by Leon Merten Lohse's avatar Leon Merten Lohse
Browse files

more elegant implementation of fftfreqn

parent 4b4db8af
......@@ -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):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment