"""
Module containing useful FFT based function and classes
"""
import numpy
[docs]def ft(data, delta):
"""
A properly scaled 1-D FFT
Parameters:
data (ndarray): An array on which to perform the FFT
delta (float): Spacing between elements
Returns:
ndarray: scaled FFT
"""
DATA = numpy.fft.fftshift(
numpy.fft.fft(
numpy.fft.fftshift(data, axes=(-1))),
axes=(-1)) * delta
return DATA
[docs]def ift(data, delta_f):
"""
Scaled inverse 1-D FFT
Parameters:
data (ndarray): Data in Fourier Space to transform
delta_f (ndarray): Frequency spacing of grid
Returns:
ndarray: Scaled data in real space
"""
DATA = numpy.fft.ifftshift(
numpy.fft.ifft(
numpy.fft.ifftshift(data, axes=(-1))),
axes=(-1)) * data.shape[-1] * delta_f
return DATA
[docs]def ft2(data, delta):
"""
A properly scaled 2-D FFT
Parameters:
data (ndarray): An array on which to perform the FFT
delta (float): Spacing between elements
Returns:
ndarray: scaled FFT
"""
DATA = numpy.fft.fftshift(
numpy.fft.fft2(
numpy.fft.fftshift(data, axes=(-1,-2))
), axes=(-1,-2)
)*delta**2
return DATA
[docs]def ift2(data, delta_f):
"""
Scaled inverse 2-D FFT
Parameters:
data (ndarray): Data in Fourier Space to transform
delta_f (ndarray): Frequency spacing of grid
Returns:
ndarray: Scaled data in real space
"""
N = data.shape[-1]
DATA = numpy.fft.ifftshift(
numpy.fft.ifft2(
numpy.fft.ifftshift(data, axes=(-1,-2)
), axes=(-1,-2))
, axes=(-1,-2)) * (N * delta_f)**2
return DATA
[docs]def rft(data, delta):
"""
A properly scaled real 1-D FFT
Parameters:
data (ndarray): An array on which to perform the FFT
delta (float): Spacing between elements
Returns:
ndarray: scaled FFT
"""
DATA = numpy.fft.fftshift(
numpy.fft.rfft(
numpy.fft.fftshift(data, axes=(-1))),
axes=(-1)) * delta
return DATA
[docs]def irft(data, delta_f):
"""
Scaled real inverse 1-D FFT
Parameters:
data (ndarray): Data in Fourier Space to transform
delta_f (ndarray): Frequency spacing of grid
Returns:
ndarray: Scaled data in real space
"""
DATA = numpy.fft.ifftshift(
numpy.fft.irfft(
numpy.fft.ifftshift(data, axes=(-1))),
axes=(-1)) * data.shape[-1] * delta_f
return DATA
[docs]def rft2(data, delta):
"""
A properly scaled, real 2-D FFT
Parameters:
data (ndarray): An array on which to perform the FFT
delta (float): Spacing between elements
Returns:
ndarray: scaled FFT
"""
DATA = numpy.fft.fftshift(
numpy.fft.rfft2(
numpy.fft.fftshift(data, axes=(-1,-2))
), axes=(-1,-2)
)*delta**2
return DATA
[docs]def irft2(data, delta_f):
"""
Scaled inverse real 2-D FFT
Parameters:
data (ndarray): Data in Fourier Space to transform
delta_f (ndarray): Frequency spacing of grid
Returns:
ndarray: Scaled data in real space
"""
N = data.shape[-1]
DATA = numpy.fft.ifftshift(
numpy.fft.irfft2(
numpy.fft.ifftshift(data, axes=(-1,-2)),
axes=(-1,-2)
),
axes=(-1,-2)) * (N * delta_f)**2
return DATA