Backends, Dealiasing & Extensions

Transfer diagnostics are configured along two orthogonal axes, plus a dealiasing strategy. This page documents each, when to use it, and how extensions are loaded.

Two backend axes

The package never conflates which transform it uses with how that work is run:

They compose: e.g. spectral = FFTSpectralBackend(), execution = ThreadedBackend() runs FFT nonlinear terms with a threaded mediator loop. A typical call:

using FFTW, OhMyThreads   # load the two extensions

result = calculate_shell_to_shell_transfer(û, ks;
    binning   = LinearBinning(1.0),
    dealiasing = OrszagTwoThirds(),
    spectral  = FFTSpectralBackend(),
    execution = ThreadedBackend())
Spectral backendDependencyBest for
DirectSumSpectralBackendNoneReference results, debugging, tiny grids (the oracle)
FFTSpectralBackendFFTWProduction spectral diagnostics on regular periodic grids
NUFFTSpectralBackendFINUFFTScattered / non-uniform Cartesian data
FSHTSpectralBackendFastSphericalHarmonicsRegular latitude–longitude grids
NUFSHTSpectralBackendNUFSHTScattered spherical observations
Execution backendDependencyBest for
SerialBackendNoneDefault; small/medium grids
ThreadedBackendOhMyThreadsMulti-core mediator/triad loops
DistributedBackendDistributed + SharedArraysMany-process single-node
GPUBackendKernelAbstractions + vendor pkgLarge grids on GPU (CUDA validated)
AutoBackendVariesAutomatic best-available execution

Dealiasing

Pseudospectral products alias high wavenumbers back onto resolved modes. Every nonlinear-term entry point takes a dealiasing::AbstractDealiasing strategy (a type, so dispatch — not a boolean — selects the path):

  • OrszagTwoThirds (default) — zero the upper third of every input field before forming the product (exact on retained modes |k| < N/3). Works with any spectral backend.
  • PaddedThreeHalves — exact 3/2 zero-padding: embed into a 3/2-sized grid, multiply, truncate back. No aliasing at all on any retained mode. FFT-only (DirectSumSpectralBackend + PaddedThreeHalves throws).
  • NoDealiasing — raw product; only for analytic single-triad tests where no aliasing can occur.
calculate_spectral_flux(û, ks; dealiasing = PaddedThreeHalves(), spectral = FFTSpectralBackend())

Aliasing breaks conservation (Σ_k T(k) ≈ 0), so the test suite asserts conservation under the dealiased paths.


Spectral backends

DirectSumSpectralBackend

The reference implementation: direct O(N²) summation for the nonlinear term. No external dependencies, always available, and the oracle every fast path is tested against.

calculate_spectral_flux(û, ks; binning = LinearBinning(1.0), spectral = DirectSumSpectralBackend())

When to use: debugging, correctness verification, very small grids (N ≤ 16).

FFTSpectralBackend

Production fast path via FFTW. Reduces the nonlinear term from O(N²) to O(Nᴰ log N): IFFT to physical space → pointwise product → FFT back → apply the dealiasing strategy. Plans are stored in the workspace and applied with mul!/ldiv!.

using FFTW   # loads the extension automatically
calculate_spectral_flux(û, ks; binning = LinearBinning(1.0), spectral = FFTSpectralBackend())

When to use: standard production runs on regular periodic grids (N ≥ 32). Also the only backend that supports PaddedThreeHalves dealiasing.

NUFFTSpectralBackend / FSHTSpectralBackend / NUFSHTSpectralBackend

Front-ends for non-uniform Cartesian (FINUFFT), regular spherical (FastSphericalHarmonics), and scattered spherical (NUFSHT) data. They transform input data to regular Fourier/spherical-harmonic coefficients, then delegate to the core spectral diagnostics.

using FastSphericalHarmonics
result = calculate_energy_transfer(
    SpectralFluxMethod(LinearBinning(1.0)),
    velocity_fields, coords, (Nθ,); spectral = FSHTSpectralBackend())

Execution backends

SerialBackend

Single-threaded, no dependencies. The default execution for every diagnostic.

ThreadedBackend

Multi-threaded via OhMyThreads with thread-local accumulators (no locks). Parallelises the outer loop over mediator shells (shell-to-shell), receiver modes (scale-to-scale), and triads (TOD).

using OhMyThreads
calculate_shell_to_shell_transfer(û, ks; binning = LinearBinning(1.0),
    spectral = FFTSpectralBackend(), execution = ThreadedBackend())
FFTW intra-transform threads ≠ `ThreadedBackend`

ThreadedBackend parallelises the outer shell/mode loop. FFTW also has its own intra-transform multithreading, set globally with FFTW.set_num_threads(n), which speeds up each individual FFT. The two are orthogonal and compose; enabling FFTW threads never changes results (asserted by the test suite). Don't oversubscribe — with execution = ThreadedBackend() each task already runs a transform, so leave FFTW at one thread (or partition cores between the two levels).

DistributedBackend

Multi-process via Distributed + SharedArrays, using @distributed (+) reduction over mediator shells / mode chunks.

using Distributed, SharedArrays
addprocs(4); @everywhere using FlowInvariantTransfer
calculate_shell_to_shell_transfer(SharedArray(û), ks;
    binning = LinearBinning(1.0), execution = DistributedBackend())

GPUBackend

Device-generic GPU execution via KernelAbstractions — custom @kernel functions for transfer density, shell accumulation (Atomix.@atomic scatter-add), and triad loops; all buffers allocated with similar(velocity_hat, …) so they follow the input array type.

using KernelAbstractions, CUDA
û_gpu = CuArray(û); ks_gpu = map(CuArray, ks)
calculate_shell_to_shell_transfer(û_gpu, ks_gpu;
    binning = LinearBinning(1.0), execution = GPUBackend(CUDABackend()))

The device kernels (transfer density for KE/helicity/enstrophy, shell reduction) and their dispatch are validated on the KernelAbstractions CPU backend (GPUBackend(KA.CPU())) against the serial reference to machine precision — so the parallel logic is correct independent of hardware. The only piece that needs a real GPU is the on-device FFT: pass spectral = FFTSpectralBackend() with a CuArray input so cuFFT rides AbstractFFTs. AMDGPU/Metal run the same kernels but are not yet hardware-validated.

using KernelAbstractions
result = calculate_shell_to_shell_transfer(û, ks;
    binning = LinearBinning(1.0), execution = GPUBackend(KA.CPU()))  # CPU backend: same kernels, no GPU

AutoBackend

Selects the best available execution backend at call time (distributed → threaded → serial). The spectral fast path is chosen independently from whichever spectral extension is loaded.


Distributed with MPI — two axes

The single-node DistributedBackend above shares one array across processes. For genuine multi-process / multi-node work there are two distinct ways to distribute, and the package provides one entry point for each (loaded by using MPI — the pencil axis also needs PencilFFTs, PencilArrays). MPI.jl bundles its own mpiexec, so a launcher works out of the box; both paths are validated single-machine with mpiexec -n 2.

Batch axis — many independent inputs (mpi_batch_map)

When you have many snapshots (a time series) that each fit in one node's memory, distribute the set of inputs: each rank applies f to a round-robin subset, then results are collated in original order (default) or reduced. No communication during each item's computation — embarrassingly parallel. This is the common post-processing mode.

using FlowInvariantTransfer, FFTW, MPI
MPI.Init()

f(û) = calculate_spectral_flux(û, ks; binning = LinearBinning(dk), spectral = FFTSpectralBackend()).flux

series = mpi_batch_map(f, snapshots)                 # Vector of per-snapshot fluxes, in order
mean_Π = mpi_batch_map(f, snapshots; reduce = :mean) # ensemble average instead of collation

reduce accepts :gather (default), :sum, :mean, or a binary combiner function; the combined result is returned on every rank.

Pencil axis — one grid too big for a node (pencil_spectral_flux)

When a single snapshot's grid doesn't fit on one node, split that grid into pencils (one slab per rank). The pseudospectral nonlinear term then needs transpose/all-to-all communication, handled by a PencilFFTs distributed FFT. The per-shell KE transfer is MPI.Allreduced to a global result identical on every rank — equal to the serial calculate_spectral_flux on the same field (validated to machine precision).

using FlowInvariantTransfer, MPI, PencilFFTs, PencilArrays
MPI.Init()

plan = build_pencil_plan((N, N), MPI.COMM_WORLD)     # auto-balanced process grid
u    = ntuple(_ -> allocate_input(plan), 2)          # fill each rank's LOCAL portion of u, v
res  = pencil_spectral_flux(u, plan, ks; binning = LinearBinning(dk))

The two axes are complementary and compose (a batch of large grids = batch axis over pencil-axis groups). The pencil path covers every invariant (KE / helicity / enstrophy) and the ShellMagnitude geometries (isotropic |k|, perpendicular k_⊥, parallel k_∥), and execution=MPIBackend(inner) sets the per-rank local backend — so MPIBackend(GPUBackend(dev)) runs a device-resident pencil (multi-GPU; the local shell reduction becomes an on-device scatter-add).


Backend Support Matrix

spectral ∈ {DirectSum, FFT}; execution ∈ {Serial, Threaded, Distributed, GPU}.

Spectral axis = which transform builds the nonlinear term (every diagnostic supports DirectSum + FFT). Execution axis = which parallelism runs the outer loop; only the diagnostics with an outer loop over shells/triads expose it (the rest run serially over the single FFT path).

DiagnosticDirectSumFFTSerialThreadedDistributedGPU
Spectral flux Π(K)
Shell-to-shell T(n,m)
Mode-to-mode S(k|p)
Smooth band-to-band T(K,Q)
Partial / decomposed fluxes
Compressible T_u(k)✓*
TOD—†

verified (serial parity to machine precision); not implemented for that axis. Coarse-graining, spherical, and the MPI batch/pencil layer are documented separately below (they are not spectral × execution diagnostics).

Notes:

  • Threaded / GPU are wired for every Fourier diagnostic. The loop methods (shell-to-shell, scale-to-scale, band-to-band, partial) thread the outer shell/mode/band/pair loop and run their per-mode transfer density through a KernelAbstractions device kernel; spectral flux threads/GPU-kernels the mode→shell reduction (device cumsum! keeps the field GPU-resident). ✓* compressible is a single FFT pipeline, so its Threaded path threads the FFTs (not an outer loop); its GPU path is a device-generic broadcast pipeline + GPUArraysCore reductions. GPU verified on GPUBackend(KA.CPU()) and on JLArrays device-generic building blocks (cuFFT rides AbstractFFTs on a CuArray by construction).
  • Distributed (Distributed + SharedArrays, many-process single-node) is implemented for spectral flux, shell-to-shell (the mode→shell scatter across workers) and TOD (the triad loop). For a grid too large for one node, use the MPI pencil axis instead (below). —† TOD's GPU path is not wired (per-triad SVD is a LAPACK, not KA-kernel, workload).
  • The fully mode-resolved S(k|p) tensor is the only query needing the O(N^{2D}) brute loop (guarded by a mode-count limit; force=true to override); T(k)/T(K,Q)/Π(K) use the fast FFT paths.
  • execution = AutoBackend() resolves to threaded when available, else serial (see resolve_backend).
  • MPI (batch + pencil axes) is a separate distribution layer (above). The pencil path supports every invariant (KE/helicity/enstrophy) and every ShellMagnitude geometry, with a 0-alloc PencilWorkspace for snapshot sweeps. Coarse-graining flux is provided by the CoarseGrainingEnergyFluxes extension (its own parallelism model); scattered coarse-graining/spherical use FINUFFT/NUFSHT (FSH for regular spherical grids).

Extension Loading

Extensions load automatically when you using their trigger package:

using FlowInvariantTransfer  # lean core only (DirectSumSpectralBackend, SerialBackend)
using FFTW                   # → FFTSpectralBackend, PaddedThreeHalves
using OhMyThreads            # → ThreadedBackend
using KernelAbstractions     # → GPUBackend (+ a vendor pkg, e.g. CUDA)
using MPI                    # → mpi_batch_map (batch axis)
using PencilFFTs, PencilArrays  # (+ MPI) → pencil_spectral_flux / build_pencil_plan (pencil axis)
using HelmholtzDecomposition # → decompose_field / Helmholtz partial fluxes
using FINUFFT                # → NUFFTSpectralBackend
using FastSphericalHarmonics # → FSHTSpectralBackend
using NUFSHT                 # → NUFSHTSpectralBackend

Calling a backend whose extension isn't loaded gives a clear error:

ArgumentError: Threaded scale-to-scale transfer requires OhMyThreads. Run `using OhMyThreads` to load the extension.