Fourier transform of a Gaussian function

import math
import numpy as np
from bokeh.io import output_notebook

import fftarray as fa
from helpers import plt_array

output_notebook(hide_banner=True)

def fa_array_assert_all_close(space: fa.Space, a: fa.Array, b: fa.Array):
    # Since we compare a real-valued and a complex-valued array, we have to give some absolute tolerance.
    # Otherwise the imaginary values which move slightly away from zero would make the comparison fail.
    np.testing.assert_allclose(a.values(space, xp=np), b.values(space, xp=np), atol=4e-15)

def gauss_pos(x, a, sigma):
    return (a * fa.exp(-(x**2/(2.* sigma**2))))/(math.sqrt(2 * np.pi) * sigma)

def gauss_freq(f, a, sigma):
    return (a * fa.exp(-(1/2)*(2*np.pi*f)**2*sigma**2))

Initialize the dimension and get the basic x and f arrays.

dim = fa.dim_from_constraints("x",
        pos_middle=1.,
        pos_extent = 10.,
        freq_middle = 2.5/(2*np.pi),
        freq_extent = 20./(2*np.pi),
        loose_params=["pos_extent", "freq_extent"]
    )


x = fa.coords_from_dim(dim, "pos")
f = fa.coords_from_dim(dim, "freq")

plt_array(x, data_name="Array (identity in pos space)")
plt_array(f, data_name="Array (identity in freq space)")

gauss_from_pos = gauss_pos(a = 1.2, x = x, sigma = 0.7)
gauss_from_freq = gauss_freq(a = 1.2, f = f, sigma = 0.7)

fa_array_assert_all_close("pos", gauss_from_pos, gauss_from_freq)
fa_array_assert_all_close("freq", gauss_from_pos, gauss_from_freq)

plt_array(gauss_from_pos, data_name="Gauss (initialised with position)")
plt_array(gauss_from_freq, data_name="Gauss (initialised with frequency)")

Just compute a centered Gaussian in position and frequency space. It does not matter where it is defined, the result is the same (up to numerical inaccuracies).

gauss_from_pos = fa.shift_freq(gauss_pos(a = 1.2, x = x, sigma = 0.7), {"x": 0.9})
gauss_from_freq = gauss_freq(a = 1.2, f = f - 0.9, sigma = 0.7)

fa_array_assert_all_close("pos", gauss_from_pos, gauss_from_freq)
fa_array_assert_all_close("freq", gauss_from_pos, gauss_from_freq)

plt_array(gauss_from_pos, data_name="Gauss (freq. shifted via FFT)")
plt_array(gauss_from_freq, data_name="Gauss (freq. shifted at init)")

Now we test this with a Gauss shifted in position space. The shift_position function multiplies a phase factor onto the Array in frequency space to achieve a user-defined shift.

gauss_from_pos = gauss_pos(a = 1.2, x = x - 0.9, sigma = 0.7)
gauss_from_freq = fa.shift_pos(gauss_freq(a = 1.2, f = f, sigma = 0.7), {"x": 0.9})

fa_array_assert_all_close("pos", gauss_from_pos, gauss_from_freq)
fa_array_assert_all_close("freq", gauss_from_pos, gauss_from_freq)

plt_array(gauss_from_pos, data_name="Gauss (pos. shifted at init)")
plt_array(gauss_from_freq, data_name="Gauss (pos. shifted via FFT)")