def fftfreq(resolution, mode='vector', dtype=None):
"""
Returns the discrete Fourier transform sample frequencies.
These are the frequencies corresponding to the components of the result of `math.fft` on a tensor of shape `resolution`.
:param resolution: grid resolution measured in cells
:param mode: one of (None, 'vector', 'absolute', 'square')
:param dtype: data type of the returned tensor
:return: tensor holding the frequencies of the corresponding values computed by math.fft
"""
assert mode in ('vector', 'absolute', 'square')
k = np.meshgrid(*[np.fft.fftfreq(int(n)) for n in resolution],
indexing='ij')
k = math.expand_dims(math.stack(k, -1), 0)
if dtype is not None:
k = k.astype(dtype)
else:
k = math.to_float(k)
if mode == 'vector':
return k
k = math.sum(k**2, axis=-1, keepdims=True)
if mode == 'square':
return k
else:
return math.sqrt(k)
def genarray(shape):
fft = randn(shape, dtype) + 1j * randn(shape, dtype)
k = fftfreq(shape[1:-1], mode='absolute')
shape_fac = math.sqrt(math.mean(shape[1:-1])) # 16: 4, 64: 8, 256: 24,
fft *= (1 / (k + 1))**power * power * shape_fac
array = math.real(math.ifft(fft)).astype(dtype)
return array
def fftfreq(resolution, mode='vector', dtype=np.float32):
assert mode in ('vector', 'absolute', 'square')
k = np.meshgrid(*[np.fft.fftfreq(int(n)) for n in resolution], indexing='ij')
k = math.expand_dims(math.stack(k, -1), 0)
k = k.astype(dtype)
if mode == 'vector':
return k
k = math.sum(k**2, axis=-1, keepdims=True)
if mode == 'square':
return k
else:
return math.sqrt(k)