Fast Fourier Transform

µGrid includes a built-in Fast Fourier Transform (FFT) engine that operates on µGrid fields. The FFT engine uses pocketfft for efficient FFT computation.

The FFTEngine class

Instantiating an FFT engine is straightforward:

import muGrid

# Create a 2D FFT engine
engine = muGrid.FFTEngine([nx, ny])

# Create a 3D FFT engine
engine = muGrid.FFTEngine([nx, ny, nz])

The FFT engine provides properties for querying the grid dimensions:

engine = muGrid.FFTEngine([16, 20])

print(f"Real-space grid: {engine.nb_domain_grid_pts}")      # (16, 20)
print(f"Fourier-space grid: {engine.nb_fourier_grid_pts}")  # (9, 20)
print(f"Backend: {engine.backend_name}")                    # "pocketfft"

The Fourier-space grid has reduced dimensions due to the real-to-complex (r2c) transform exploiting Hermitian symmetry.

Creating FFT fields

The FFT engine manages two field collections: one for real-space fields and one for Fourier-space fields. Use the engine’s methods to create fields:

import muGrid

engine = muGrid.FFTEngine([16, 20])

# Create real-space field (real-valued)
real_field = engine.real_space_field("real")
print(f"Real field shape: {real_field.s.shape}")  # (1, 1, 16, 20)

# Create Fourier-space field (complex-valued)
fourier_field = engine.fourier_space_field("fourier")
print(f"Fourier field shape: {fourier_field.s.shape}")  # (1, 1, 9, 20)

Forward and inverse transforms

The fft method performs the forward transform (real to Fourier space) and ifft performs the inverse transform (Fourier to real space):

import numpy as np
import muGrid

engine = muGrid.FFTEngine([16, 20])

real_field = engine.real_space_field("real")
fourier_field = engine.fourier_space_field("fourier")

# Initialize real-space field
real_field.s[:] = np.random.randn(1, 1, 16, 20)
original = real_field.s.copy()

# Forward FFT: real -> Fourier
engine.fft(real_field, fourier_field)

# Inverse FFT: Fourier -> real
engine.ifft(fourier_field, real_field)

# Apply normalization for roundtrip
real_field.s[:] *= engine.normalisation

# Verify roundtrip
np.testing.assert_allclose(real_field.s, original, atol=1e-14)

Normalization

µGrid FFT transforms are not normalized by default. A forward-inverse roundtrip picks up a global factor equal to the total number of grid points. This factor is available as the normalisation property:

engine = muGrid.FFTEngine([8, 10])
print(f"Normalization factor: {engine.normalisation}")  # 1/80 = 0.0125

To get the original values after a roundtrip, multiply by the normalization factor after the inverse transform.

The reason for leaving normalization to the user is that there are multiple valid conventions (normalize forward, normalize inverse, or split between both), and different applications may prefer different choices.

Frequency and coordinate arrays

The FFT engine provides properties for accessing frequency and coordinate arrays that are correctly shaped for the local subdomain. This is particularly useful for MPI-parallel computations where each rank only handles a portion of the grid.

Frequency arrays

The fftfreq and ifftfreq properties provide FFT frequency arrays:

import muGrid
import numpy as np

engine = muGrid.FFTEngine([7, 4])

# Normalized frequencies (range [-0.5, 0.5))
qx, qy = engine.fftfreq
print(f"fftfreq shape: {engine.fftfreq.shape}")  # (2, 4, 4) for r2c transform

# Integer frequency indices
iqx, iqy = engine.ifftfreq
print(f"ifftfreq dtype: {iqx.dtype}")  # integer type

# These match numpy's frequency arrays (sliced for r2c transform)
freq_ref = np.array(np.meshgrid(
    *(np.fft.fftfreq(n) for n in [7, 4]), indexing="ij"
))
freq_ref = freq_ref[:, :4, :]  # Slice for half-complex
np.testing.assert_allclose(engine.fftfreq, freq_ref)

The frequency arrays have shape [dim, local_fx, local_fy, ...] where the first axis indexes the spatial dimension and the remaining axes match the local Fourier subdomain dimensions.

Coordinate arrays

The coords and icoords properties provide real-space coordinate arrays:

import muGrid

engine = muGrid.FFTEngine([7, 4])

# Normalized coordinates (range [0, 1))
x, y = engine.coords
print(f"coords shape: {engine.coords.shape}")  # (2, 7, 4)

# Integer coordinate indices
ix, iy = engine.icoords
print(f"icoords dtype: {ix.dtype}")  # integer type

# Verify: integer coords = fractional * n
np.testing.assert_allclose(ix, x * 7)
np.testing.assert_allclose(iy, y * 4)

For fields with ghost regions, use coordsg and icoordsg to get coordinates including the ghost cells:

engine = muGrid.FFTEngine(
    [64, 64],
    nb_ghosts_left=[1, 1],
    nb_ghosts_right=[1, 1]
)

# Coordinates without ghosts
x, y = engine.coords

# Coordinates with ghosts (larger array)
xg, yg = engine.coordsg
print(f"Without ghosts: {engine.coords.shape}")   # (2, local_nx, local_ny)
print(f"With ghosts: {engine.coordsg.shape}")     # (2, local_nx+2, local_ny+2)

MPI-parallel frequency arrays

In MPI-parallel runs, frequency and coordinate arrays return only the data for the local subdomain owned by each rank:

from mpi4py import MPI
import muGrid

comm = muGrid.Communicator(MPI.COMM_WORLD)
engine = muGrid.FFTEngine([128, 128], comm)

# Each rank gets frequencies for its local Fourier subdomain
qx, qy = engine.fftfreq
print(f"Rank {MPI.COMM_WORLD.rank}: fftfreq shape = {engine.fftfreq.shape}")

# Coordinates for local real-space subdomain
x, y = engine.coords
print(f"Rank {MPI.COMM_WORLD.rank}: coords shape = {engine.coords.shape}")

Example: Fourier-space operations

A common pattern is using frequency arrays to construct Fourier-space operators:

import numpy as np
import muGrid

engine = muGrid.FFTEngine([64, 64])
qx, qy = engine.fftfreq

real_field = engine.real_space_field("u")
fourier_field = engine.fourier_space_field("u_hat")

# Initialize with a cosine wave
x, y = engine.coords
real_field.p[0] = np.cos(2 * np.pi * x)

# Forward FFT
engine.fft(real_field, fourier_field)

# Compute Laplacian in Fourier space: -4π²(qx² + qy²) * u_hat
laplacian_hat = -4 * np.pi**2 * (qx**2 + qy**2) * fourier_field.p[0]

# Or use integer frequencies for identifying specific modes
iqx, iqy = engine.ifftfreq
# Find the (1, 0) mode
mode_mask = (np.abs(iqx) == 1) & (iqy == 0)

Hermitian grid dimensions

For real-to-complex transforms, the output has reduced dimensions due to Hermitian symmetry. Use get_hermitian_grid_pts to compute the Fourier-space grid size:

import muGrid

real_grid = [64, 64, 64]
fourier_grid = muGrid.get_hermitian_grid_pts(real_grid)
print(f"Fourier grid: {fourier_grid}")  # [33, 64, 64]

The first dimension is reduced to n//2 + 1.

FFT normalization factor

The fft_normalization function returns the normalization factor for a given grid size:

import muGrid

norm = muGrid.fft_normalization([16, 20])
print(f"Normalization: {norm}")  # 1/320

Example: Fourier derivative

A common use of FFTs is computing derivatives in Fourier space. Here’s an example computing the gradient of a 2D field:

import numpy as np

import muGrid

# Grid parameters
nx, ny = 64, 64
Lx, Ly = 2 * np.pi, 2 * np.pi

engine = muGrid.FFTEngine([nx, ny])

# Create fields
f = engine.real_space_field("f")
f_hat = engine.fourier_space_field("f_hat")
grad_x = engine.real_space_field("grad_x")
grad_y = engine.real_space_field("grad_y")
grad_hat = engine.fourier_space_field("grad_hat")

# Get real-space coordinates from engine
X, Y = engine.coords
X = X * Lx  # Scale from [0, 1) to [0, Lx)
Y = Y * Ly

# Initialize with a smooth function: f = sin(x) * cos(y)
f.s[0, 0, :, :] = np.sin(X) * np.cos(Y)

# Forward FFT
engine.fft(f, f_hat)

# Get wavevectors from engine (integer indices, then scale)
# fftfreq returns normalized frequencies in [-0.5, 0.5)
# Multiply by 2*pi*n/L to get physical wavevectors
QX, QY = engine.fftfreq
KX = 2 * np.pi * QX * nx / Lx  # Physical wavevector
KY = 2 * np.pi * QY * ny / Ly

# Derivative in x: multiply by i*kx
grad_hat.s[0, 0, :, :] = 1j * KX * f_hat.s[0, 0, :, :]
engine.ifft(grad_hat, grad_x)
grad_x.s[:] *= engine.normalisation

# Derivative in y: multiply by i*ky
grad_hat.s[0, 0, :, :] = 1j * KY * f_hat.s[0, 0, :, :]
engine.ifft(grad_hat, grad_y)
grad_y.s[:] *= engine.normalisation

# Verify: d/dx[sin(x)cos(y)] = cos(x)cos(y)
expected_grad_x = np.cos(X) * np.cos(Y)
np.testing.assert_allclose(grad_x.s[0, 0], expected_grad_x, atol=1e-10)

print("Fourier derivative test passed!")

Multi-component fields

FFT fields can have multiple components:

import muGrid

engine = muGrid.FFTEngine([16, 20])

# Create field with 3 components (e.g., velocity vector)
velocity = engine.real_space_field("velocity", nb_components=3)
velocity_hat = engine.fourier_space_field("velocity_hat", nb_components=3)

print(f"Velocity shape: {velocity.s.shape}")      # (3, 1, 16, 20)
print(f"Velocity_hat shape: {velocity_hat.s.shape}")  # (3, 1, 9, 20)

# FFT transforms all components - fields are passed directly
engine.fft(velocity, velocity_hat)
engine.ifft(velocity_hat, velocity)

MPI-parallel FFT

The FFTEngine class supports MPI parallelization using pencil (2D) decomposition, which allows efficient scaling to large numbers of ranks.

Basic parallel usage

import numpy as np
from mpi4py import MPI
import muGrid
from muGrid import Communicator

# Create parallel FFT engine
nb_grid_pts = [128, 128, 128]
comm = Communicator(MPI.COMM_WORLD)
engine = muGrid.FFTEngine(nb_grid_pts, comm)

# Each rank has a subdomain
print(f"Rank {MPI.COMM_WORLD.rank}:")
print(f"  Global grid: {engine.nb_domain_grid_pts}")
print(f"  Local subdomain: {engine.nb_subdomain_grid_pts}")
print(f"  Process grid: {engine.process_grid}")
print(f"  Process coords: {engine.process_coords}")

The FFT engine uses pencil decomposition, distributing the grid across a 2D process grid for efficient scaling.

FFT with ghost regions

For computations requiring ghost cells (e.g., stencil operations), specify the ghost buffer sizes:

import muGrid
from muGrid import Communicator
from mpi4py import MPI

comm = Communicator(MPI.COMM_WORLD)

# Create FFT engine with ghost regions
engine = muGrid.FFTEngine(
    nb_domain_grid_pts=[64, 64, 64],
    comm=comm,
    nb_ghosts_left=[1, 1, 1],
    nb_ghosts_right=[1, 1, 1]
)

# Real-space fields include ghost regions
real_field = engine.real_space_field("displacement", nb_components=3)

# Fourier-space fields have no ghosts (hard assumption)
fourier_field = engine.fourier_space_field("displacement_k", nb_components=3)

Run with MPI:

$ mpirun -np 4 python parallel_fft.py

GPU support

When µGrid is compiled with CUDA or HIP support, FFT operations can run on GPU. See the GPU Computing documentation for details on building with GPU support and working with GPU fields.

Currently, the FFT engine operates on host (CPU) memory. For GPU-accelerated FFT, you can transfer data between host and device fields as needed.