FFT
Parallel fast Fourier transforms with amrex::FFT and the FFT-based
Poisson solvers built on them. The transforms work directly on
MultiFab data distributed over MPI ranks. A CPU build uses FFTW; GPU
builds use cuFFT, rocFFT or oneMKL. These tutorials require AMReX built
with AMReX_FFT=ON (the default in the superbuild) and run in 2D and 3D.
Build only this topic with -DAMReXTutorials_TOPICS=FFT; see
Getting Started.
Basic
- Level:
★
- Source:
- Key APIs:
FFT::R2C,getSpectralDataLayout,forward,backward,forwardThenBackward,scalingFactor,cMultiFab,GpuComplex
Real-to-complex transforms of a MultiFab on a periodic unit domain: a
forward and backward round trip, a spectral derivative \(d/dx\), and a
low-pass filter that zeroes all modes with \(|k| > k_\mathrm{cut}\).
The test field is a smooth low-wavenumber part plus a high-wavenumber part
that the filter removes, and all three results are checked against the exact
answer to round-off.
The spectral data has \((n_x/2+1, n_y, n_z)\) points: only
\(k_x \ge 0\) is stored, because the transform of real data is
Hermitian. In \(y\) and \(z\), indices above \(n/2\) are
negative wavenumbers. FFTs are unnormalized, so a round trip multiplies the
data by the number of points; scalingFactor() returns its inverse.
An R2C object does the expensive setup once and is reused. The spectral
data can be stored in a cMultiFab with the layout given by
getSpectralDataLayout:
// R2C does the expensive setup (plans, communication patterns) once;
// keep it around and reuse it.
FFT::R2C<Real,FFT::Direction::both> r2c(geom.Domain());
// The real-to-complex transform stores only half of the spectrum in x:
// (nx/2+1, ny, nz) complex values, distributed over the ranks.
auto const& [cba, cdm] = r2c.getSpectralDataLayout();
cMultiFab phi_hat(cba, cdm, 1, 0);
amrex::Print() << "Real domain " << geom.Domain() << "\n"
<< "Spectral domain " << cba.minimalBox() << "\n";
// 1. Round trip. FFTs are unnormalized, so forward followed by
// backward multiplies the data by the number of points.
MultiFab phi2(ba, dm, 1, 0);
r2c.forward(phi, phi_hat);
r2c.backward(phi_hat, phi2);
phi2.mult(r2c.scalingFactor()); // 1/(nx*ny*nz)
Real err_roundtrip = max_diff(phi, phi2);
amrex::Print() << "Round trip error: " << err_roundtrip << "\n";
forwardThenBackward calls a lambda at every spectral point between the
two transforms, without copying the spectral data into a cMultiFab.
The spectral derivative multiplies by \(i k_x\):
// 2. Spectral derivative d/dx: multiply by i*kx. forwardThenBackward
// calls the lambda at every spectral point (i,j,k); it avoids the
// communication of copying the spectral data to a cMultiFab.
MultiFab dphi(ba, dm, 1, 0);
r2c.forwardThenBackward(phi, dphi,
[=] AMREX_GPU_DEVICE (int i, int, int, GpuComplex<Real>& sp)
{
// i is in [0,nx/2]. Zero the Nyquist mode, it has no sign.
Real kx = (2*i == nx) ? Real(0) : Real(2)*Math::pi<Real>()*Real(i);
sp *= GpuComplex<Real>(Real(0), kx*scale);
});
The low-pass filter maps the indices \(j\) and \(k\) to signed wavenumbers:
// 3. Low-pass filter: zero all modes with |k| > kcut. In y and z, indices
// above n/2 are negative wavenumbers.
MultiFab phi_low(ba, dm, 1, 0);
r2c.forwardThenBackward(phi, phi_low,
[=] AMREX_GPU_DEVICE (int i, int j, int k, GpuComplex<Real>& sp)
{
int ky = (j <= ny/2) ? j : j-ny;
int kz = (k <= nz/2) ? k : k-nz;
int k2 = i*i + ky*ky + kz*kz;
sp *= (k2 > kcut*kcut) ? Real(0) : scale;
});
Run (from the top of the repository, see Getting Started):
cmake --build build --target FFT_Basic -j 8
cd build/FFT/Basic
mpiexec -n 2 ./Basic inputs
Runtime parameters: n_cell, max_grid_size and kcut.
Poisson
- Level:
★★
- Source:
- Key APIs:
FFT::Poisson,FFT::PoissonHybrid,FFT::Boundary
Solve \(\nabla^2\phi = f\) with FFT-based direct solvers and check the result against a manufactured solution:
FFT::Poissonwith periodic boundaries.FFT::Poissonwith homogeneous Dirichlet (FFT::Boundary::odd, sine transform) and Neumann (FFT::Boundary::even, cosine transform) boundaries, set per face.3D only:
FFT::PoissonHybrid, which uses FFTs in \(x\) and \(y\) and a tridiagonal solve in \(z\). The \(z\) boundaries must be Dirichlet or Neumann, and \(z\) may have a non-uniform cell size.
The solvers invert the standard second-order Laplacian exactly, so the
residual of the discrete equations is at round-off level, while the error
against the exact solution is \(O(\Delta x^2)\). If phi has ghost
cells, the solver fills one layer from the boundary conditions. With all
boundaries periodic or Neumann, the solution is unique only up to a constant
and the right-hand side must have zero mean.
The boundary types are given as a pair of FFT::Boundary values (low and
high face) per direction, with a matching Geometry periodicity:
// 2. Homogeneous Dirichlet (odd) and Neumann (even) boundaries, using
// sine and cosine transforms: x Dirichlet on both sides, y Neumann on
// both sides, z Dirichlet at zlo and Neumann at zhi.
{
Array<BCPair,AMREX_SPACEDIM> bc{AMREX_D_DECL(
std::make_pair(FFT::Boundary::odd, FFT::Boundary::odd),
std::make_pair(FFT::Boundary::even, FFT::Boundary::even),
std::make_pair(FFT::Boundary::odd, FFT::Boundary::even))};
Geometry geom(domain, rb, CoordSys::cartesian, {AMREX_D_DECL(0,0,0)});
FFT::Poisson<MultiFab> solver(geom, bc);
Exact exact{{AMREX_D_DECL(Real(1),Real(1),Real(0.5))},
{AMREX_D_DECL(0,1,0)}};
pass = run_case("Dirichlet/Neumann (FFT::Poisson)", solver, geom, ba, dm, exact)
&& pass;
}
FFT::PoissonHybrid is set up in the same way:
// 3. FFT in x and y (periodic here), tridiagonal solve in z (Neumann at
// zlo, Dirichlet at zhi). PoissonHybrid also accepts a non-uniform
// dz: solver.solve(phi, rhs, dz).
{
Array<BCPair,AMREX_SPACEDIM> bc{
BCPair{FFT::Boundary::periodic, FFT::Boundary::periodic},
BCPair{FFT::Boundary::periodic, FFT::Boundary::periodic},
BCPair{FFT::Boundary::even, FFT::Boundary::odd}};
Geometry geom(domain, rb, CoordSys::cartesian, {1,1,0});
FFT::PoissonHybrid<MultiFab> solver(geom, bc);
Exact exact{{Real(2),Real(2),Real(0.5)}, {0,0,1}};
pass = run_case("Periodic x,y + Neumann/Dirichlet z (FFT::PoissonHybrid)",
solver, geom, ba, dm, exact) && pass;
}
Run (from the top of the repository, see Getting Started):
cmake --build build --target FFT_Poisson -j 8
cd build/FFT/Poisson
mpiexec -n 2 ./Poisson inputs
Doubling n_cell reduces the error by about a factor of four.
OpenBC
- Level:
★★
- Source:
- Key APIs:
FFT::OpenBCSolver,setGreensFunction,FFT::PoissonOpenBC
Solve \(\nabla^2\phi = \rho\) in free space (open boundaries) for a Gaussian charge and compare with the analytic potential. In 3D,
and in 2D, \(\phi = \left(\ln r + E_1\left(r^2/(2\sigma^2)\right)/2\right)/(2\pi)\), where \(E_1\) is the exponential integral.
The free-space solution is the convolution of \(\rho\) with the Green’s
function \(G\), which is \(-1/(4\pi r)\) in 3D and
\(\ln(r)/(2\pi)\) in 2D. FFT::OpenBCSolver computes it with FFTs on
a domain doubled in each direction (Hockney’s method), so that periodic
images do not interact. The user provides \(G\) integrated over a cell
with setGreensFunction; here with a \(2^D\)-point Gauss quadrature,
which also avoids the singularity at \(r = 0\):
// The open-BC solution is the convolution phi = G * rho with the
// free-space Green's function G. OpenBCSolver computes it with FFTs on
// a domain doubled in each direction (Hockney's method). The user
// provides G integrated over a cell, here with 2^D-point Gauss
// quadrature, which also avoids the singularity at r = 0.
FFT::OpenBCSolver<Real> solver(geom.Domain());
{
auto const lo = geom.Domain().smallEnd();
Real const gfac = Real(1)/std::sqrt(Real(12));
#if (AMREX_SPACEDIM == 3)
Real const fac = -dx[0]*dx[1]*dx[2] / (Real(8) * Real(4)*Math::pi<Real>());
#else
Real const fac = dx[0]*dx[1] / (Real(4) * Real(2)*Math::pi<Real>());
#endif
// (i,j,k) - lo is the separation between two cells in cell units.
solver.setGreensFunction([=] AMREX_GPU_DEVICE (int i, int j, int k) -> Real
{
amrex::ignore_unused(k); // unused in 2D
IntVect const iv(AMREX_D_DECL(i,j,k));
Real r = 0;
#if (AMREX_SPACEDIM == 3)
for (int gz = -1; gz <= 1; gz += 2) {
#endif
for (int gy = -1; gy <= 1; gy += 2) {
for (int gx = -1; gx <= 1; gx += 2) {
IntVect const g(AMREX_D_DECL(gx,gy,gz));
Real d2 = 0;
for (int d = 0; d < AMREX_SPACEDIM; ++d) {
Real xg = (Real(iv[d]-lo[d]) + Real(g[d])*gfac) * dx[d];
d2 += xg*xg;
}
#if (AMREX_SPACEDIM == 3)
r += Real(1)/std::sqrt(d2); // G = -1/(4 pi r)
#else
r += Real(0.5)*std::log(d2); // G = ln(r)/(2 pi)
#endif
}}
#if (AMREX_SPACEDIM == 3)
}
#endif
return fac * r;
});
}
MultiFab phi(ba, dm, 1, 0);
solver.solve(phi, rho);
In 3D, FFT::PoissonOpenBC provides the same Green’s function:
FFT::PoissonOpenBC<MultiFab> poisson(geom);
MultiFab phi2(ba, dm, 1, 0);
poisson.solve(phi2, rho);
The error is second order in \(\Delta x/\sigma\) and is checked against a tolerance.
Run (from the top of the repository, see Getting Started):
cmake --build build --target FFT_OpenBC -j 8
cd build/FFT/OpenBC
mpiexec -n 2 ./OpenBC inputs
Runtime parameters: n_cell, max_grid_size and the charge width
sigma. With n_cell = 64, the relative error drops from about 0.007
to 0.0018.
EnergySpectrum
- Level:
★★
- Source:
- Key APIs:
FFT::R2C,FFT::Info::setBatchSize,cMultiFab,HostDevice::Atomic::Add,Gpu::DeviceVector,Gpu::copy,ParallelDescriptor::ReduceRealSum
Compute the shell-averaged kinetic energy spectrum \(E(k)\) of a periodic velocity field, a standard diagnostic for turbulence. The velocity field is synthetic: three Fourier modes with \(|k| = 2, 5, 8\) (in units of \(2\pi/L\)) and known energies. The program checks that the energy is in the expected shells and that Parseval’s theorem holds: \(\sum_k E(k) = \langle|\mathbf{u}|^2\rangle/2\).
All velocity components are transformed in one batched call:
// Batched forward transform: all velocity components in one call.
FFT::Info info{};
info.setBatchSize(ncomp);
FFT::R2C<Real,FFT::Direction::forward> r2c(geom.Domain(), info);
auto const& [cba, cdm] = r2c.getSpectralDataLayout();
cMultiFab vel_hat(cba, cdm, ncomp, 0);
r2c.forward(vel, vel_hat);
Each spectral point adds \(|\hat{u}|^2/(2N^2)\) to the shell \(\mathrm{round}(|k|)\) with an atomic add, giving a histogram on the device. Because only \(k_x \ge 0\) is stored, points with \(0 < k_x < n_x/2\) count twice. The local histograms are then summed across MPI ranks:
// Histogram on the device: each spectral point adds its energy to
// its shell with an atomic add.
Gpu::DeviceVector<Real> ek_d(nshells, Real(0));
Real* pek = ek_d.data();
auto const& vh = vel_hat.const_arrays();
ParallelFor(vel_hat, [=] AMREX_GPU_DEVICE (int b, int i, int j, int k)
{
// Only kx >= 0 is stored. Indices above n/2 are negative wavenumbers.
int ky = (j <= ny/2) ? j : j-ny;
int kz = (k <= nz/2) ? k : k-nz;
int shell = int(std::round(std::sqrt(Real(i*i + ky*ky + kz*kz))));
Real e = 0;
for (int n = 0; n < ncomp; ++n) {
e += amrex::norm(vh[b](i,j,k,n)); // |u_hat|^2
}
// Hermitian symmetry: u_hat(-k) = conj(u_hat(k)). Except for kx = 0
// and the Nyquist mode, each stored point stands for two modes.
if (i > 0 && 2*i != nx) { e *= Real(2); }
HostDevice::Atomic::Add(pek+shell, e*fac);
});
// Bring the local sums to the host and add them across ranks.
std::vector<Real> ek(nshells);
Gpu::copy(Gpu::deviceToHost, ek_d.begin(), ek_d.end(), ek.begin());
ParallelDescriptor::ReduceRealSum(ek.data(), nshells);
Run (from the top of the repository, see Getting Started):
cmake --build build --target FFT_EnergySpectrum -j 8
cd build/FFT/EnergySpectrum
mpiexec -n 2 ./EnergySpectrum inputs
The spectrum is written to spectrum.txt (set with outfile) in two
columns, \(k\) and \(E(k)\), and the program prints:
Nonzero shells:
E(2) = 0.25 (exact 0.25)
E(5) = 0.0625 (exact 0.0625)
E(8) = 0.015625 (exact 0.015625)
Kinetic energy: physical space 0.328125, spectral space 0.328125