Multi-dimensional FFT

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 gauss_pos(x, a, sigma):
    return (a * np.exp(-(x**2/(2.* sigma**2))))/(np.sqrt(2 * np.pi) * sigma)

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

Define two Dimension objects with different names “x” and “y”.

x_dim = fa.dim_from_constraints(
        name="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"]
    )

y_dim = fa.dim_from_constraints(
        name="y",
        pos_middle=0.,
        pos_extent = 5.,
        freq_middle = 0,
        n=64,
    )

x_dim_array = fa.coords_from_dim(x_dim, "pos")
y_dim_array = fa.coords_from_dim(y_dim, "pos")

plt_array(x_dim_array, data_name="Array with x_dim")
plt_array(y_dim_array, data_name="Array with y_dim")

While one Dimension object always reflects the properties of exactly one dimension (here: “x” or “y”), an Array object can also be multi-dimensional.

We can create a multi-dimensional Array object by a mathematical operation between Array objects with different Dimensions, i.e., different names. Here, we will add two different Arrays which are combined by broadcasting along different dimensions.

# Array objects are broadcasted along different dimensions
array_2d = x_dim_array + y_dim_array
# You can also combine Array objects with only partially shared dimensions
array_still_2d = array_2d + x_dim_array

print("2d Array dimensions:", array_2d.dims_dict)
print("Dimension check:", array_2d.dims_dict == array_still_2d.dims_dict)

2d Array dimensions: {'x': Dimension(name='x', n=64, d_pos=0.20979344407304665, pos_min=-5.713390210337493, freq_min=-1.985409232920197, dynamically_traced_coords=False), 'y': Dimension(name='y', n=64, d_pos=0.07936507936507936, pos_min=-2.5396825396825395, freq_min=-6.3, dynamically_traced_coords=False)}
Dimension check: True

A 2d Array can be transformed between position and frequency space in each dimension, separately. Here, we choose to represent the Array in always the same space for both dimension. Below, we show its complex values along all points in the 2d “xy” space.

plt_array(array_2d, data_name="2d Array")