import numpy as np
from scipy.fft import fft, ifft, rfft, irfft
[docs]def my_fft(u, n=None, complex=False):
r"""Takes FFT of a real or complex array along the -1 axis, optimizing storage if the array is real.
Wraps the scipy.fft routines fft, rfft (see https://docs.scipy.org/doc/scipy/reference/fft.html).
Parameters
----------
u : ndarray
Real or complex numpy array.
n: int or None, optional.
Total number of frequencies to include, if padding. Default: None (no padding).
complex : boolean, optional
Equal to True if u is complex and False if u is real. Default: False.
Returns
-------
out : complex ndarray
Same shape conventions as scipy.fft.fft and scipy.fft.rfft (see https://docs.scipy.org/doc/scipy/reference/fft.html).
"""
if complex:
out = fft(u, n=n)
else:
out = rfft(u, n=n)
return out
[docs]def my_ifft(V, n=None, complex=False):
r"""Takes inverse FFT of an array along the -1 axis, optimizing storage if the inverse FFT is expected to be real.
Wraps the scipy.fft routines ifft, irfft (see https://docs.scipy.org/doc/scipy/reference/fft.html).
Parameters
----------
V : complex ndarray
Array we want to apply inverse FFT to.
n: int or None, optional.
Total number of frequencies to pad V with, if padding is desired. Default: None (no padding).
complex : boolean, optional
Equal to True if inverse FFT of V is known to be complex and False if inverse FFT is known to be
real. Default: False.
Returns
-------
out: ndarray
Same shape conventions as scipy.fft.ifft and scipy.fft.irfft (see https://docs.scipy.org/doc/scipy/reference/fft.html).
"""
if complex:
out = ifft(V, n=n)
else:
out = irfft(V, n=n)
return out
[docs]def integrate(u, length):
r"""Integrates N samples of a real or complex space-time field over the spatial interval [-0.5*length, 0.5*length] using the FFT.
Parameters
----------
u : ndarray
Real or complex array to be integrated in space. Convention is that spatial variation is stored in the
-1 axis.
length : float
Total length of spatial domain.
Returns
-------
out: float or ndarray
Approximation of the integral of the sampled space-time field over the spatial interval
[-0.5*length, 0.5*length].
"""
N = np.shape(u)[-1]
return (length/N) * np.real(fft(u, axis=-1)[..., 0])
[docs]def dealiased_pow(V,p):
r"""Compute the Fourier-space version of a power of an array, dealiased by zero-padding.
Important: this only works for real-valued arrays because for complex arrays :math:`u`, algebraic nonlinearities
typically involve the complex conjugate :math:`\overline{u}` as well as :math:`u`. So, padding for such
nonlinearities terms should be done on-the-fly.
A future release of *joe* may include support for dealiased complex powers.
Parameters
----------
V : complex ndarray
p : int
Returns
-------
out: complex ndarray
A version of :math:`\mathcal{F}\left[\left(\mathcal{F}^{-1}V\right)^p\right]` that has been dealiased
via zero-padding.
"""
#
N = 2 * (len(V) - 1)
K = int(0.5*(p+1)*N)
upad = my_ifft(V, n=K, complex=False)
# FOR CLARITY/DEV. PURPOSES ONLY: the above lines of code produce the same result as the following block:
#Vpad = 1j * np.zeros(int(0.5 * K + 1), dtype=float)
#Vpad[0:len(V)] = V
#upad = irfft(Vpad, n=K)
out = ((K/N)**(p-1))*my_fft(upad ** p, complex=False)[0:int(0.5*N+1)]
# TODO: I discovered the correct normalizations via trial-and error, and by comparing Fu2, Fu3 against
# rfft(irfft(V) ** 2), rfft(irfft(V) ** 3) resp. Make sure to write in a tutorial, or in the docs, a more
# systematic way of determining this normalization. Dealiasing via padding would be a great topic for a
# "tutorial 5"! U could also experiment with filtering vs. padding vs. doing nothing!
return out