PoPS/doxygen/poisson__fft__solver_8hpp_source.html

#pragma once


#include <pops/numerics/elliptic/interface/elliptic_solver.hpp>

#include <pops/numerics/elliptic/poisson/poisson_fft.hpp>

#include <pops/numerics/elliptic/poisson/poisson_operator.hpp>

#include <pops/mesh/layout/box_array.hpp>

#include <pops/mesh/layout/distribution_mapping.hpp>

#include <pops/mesh/boundary/fill_boundary.hpp>

#include <pops/mesh/geometry/geometry.hpp>

#include <pops/mesh/storage/mf_arith.hpp>

#include <pops/mesh/storage/multifab.hpp>

#include <pops/mesh/boundary/physical_bc.hpp>

#include <pops/parallel/comm.hpp>


#include <cassert>

#include <cstddef>

#include <functional>

#include <stdexcept>

#include <vector>


#ifdef POPS_HAS_MPI

#include <mpi.h>

#endif


namespace pops {


class PoissonFFTSolver {

 public:


  PoissonFFTSolver(const Geometry& geom, const BoxArray& ba, const BCRec& = BCRec{},

                   std::function<bool(Real, Real)> = {}, bool spectral = false)

      : geom_(geom),

        dm_(ba.size(), n_ranks()),

        phi_(ba, dm_, 1, 1),

        rhs_(ba, dm_, 1, 0),

        res_(ba, dm_, 1, 0),

        fft_(geom.domain.nx(), geom.domain.ny(), geom.xhi - geom.xlo, geom.yhi - geom.ylo,

             spectral) {

    // HARD guard (active in Release, NDEBUG does NOT remove it): this direct solver is single-rank /

    // single box. Under a system DistributionMapping with n_ranks()>1, some ranks have no local

    // box (local_size()==0) and solve() would dereference a non-existent fab(0) -> SIGSEGV. The old

    // assert vanished in Release and the protection silently disappeared. We throw on ALL ranks

    // (each one constructs the object), so no deadlock. For the distributed periodic case:

    // DistributedFFTSolver (slabs, MPI_Alltoall).

    if (n_ranks() != 1)

      throw std::runtime_error(

          "fft solver unsupported in MPI (n_ranks>1): use geometric_mg or the distributed fft "

          "solver (DistributedFFTSolver)");

    if (ba.size() != 1)

      throw std::runtime_error(

          "PoissonFFTSolver: single box required (ba.size()==1); for a distributed multi-box "

          "domain, use DistributedFFTSolver or geometric_mg");

  }


  MultiFab& rhs() { return rhs_; }

  MultiFab& phi() { return phi_; }

  const Geometry& geom() const { return geom_; }


  // Solve lap(phi) = rhs in place (direct, one forward FFT + inverse).


  void solve() {

    const int Nx = geom_.domain.nx(), Ny = geom_.domain.ny();

    const ConstArray4 f = rhs_.fab(0).const_array();

    const Box2D v = rhs_.box(0);

    std::vector<double> rho(static_cast<std::size_t>(Nx) * Ny), phil;

    for (int j = 0; j < Ny; ++j)

      for (int i = 0; i < Nx; ++i)

        rho[static_cast<std::size_t>(j) * Nx + i] = f(v.lo[0] + i, v.lo[1] + j);

    fft_.solve(rho, phil);  // mode k=0 (constant) set to zero -> phi with zero mean

    Array4 p = phi_.fab(0).array();

    for (int j = 0; j < Ny; ++j)

      for (int i = 0; i < Nx; ++i)

        p(v.lo[0] + i, v.lo[1] + j) = phil[static_cast<std::size_t>(j) * Nx + i];

    // phi GHOSTS: the aux derivation (centered grad phi, solve_fields) reads the i+-1/j+-1

    // neighbors, hence the domain-boundary ghosts. GeometricMG fills them while smoothing; this

    // direct solver writes ONLY the valid cells -> without this exchange, the boundary gradient

    // would read stale ghosts (latent bug exposed by an electric source wired onto 'fft':

    // wrong Ex on the boundary ring). Pure periodic by construction (ctor guards).

    fill_boundary(phi_, geom_.domain, Periodicity{true, true});

  }


  // Discrete residual ||lap(phi) - rhs|| (~ round-off for this direct solver).


  Real residual() {

    BCRec bc;  // periodic

    poisson_residual(phi_, rhs_, geom_, bc, res_);

    return norm_inf(res_);

  }


 private:

  Geometry geom_;

  DistributionMapping dm_;

  MultiFab phi_, rhs_, res_;

  PoissonFFT fft_;

};


static_assert(EllipticSolver<PoissonFFTSolver>, "PoissonFFTSolver must model EllipticSolver");


class DistributedFFTSolver {

 public:


  DistributedFFTSolver(const Geometry& geom, const BCRec& = BCRec{},

                       std::function<bool(Real, Real)> = {})

      : geom_(geom),

        Nx_(geom.domain.nx()),

        nyl_(geom.domain.ny() / n_ranks()),

        fft_(geom.domain.nx(), geom.domain.ny(), geom.xhi - geom.xlo, geom.yhi - geom.ylo) {

    const int np = n_ranks();

    // HARD guard (active in Release): nyl_ = Ny / np is already computed in the init list.

    // If Ny is not divisible by np, the slabs would be mis-sized and solve() would read out

    // of bounds. The old assert vanished under NDEBUG. Aligned with PoissonFFTSolver (throw).

    if (geom.domain.ny() % np != 0)

      throw std::runtime_error("DistributedFFTSolver: Ny must be divisible by n_ranks()");

    std::vector<Box2D> slabs;

    for (int r = 0; r < np; ++r)

      slabs.push_back(Box2D{{0, r * nyl_}, {Nx_ - 1, (r + 1) * nyl_ - 1}});

    BoxArray ba(std::move(slabs));

    dm_ = DistributionMapping(np, np);  // slab r -> rank r

    phi_ = MultiFab(ba, dm_, 1, 1);

    rhs_ = MultiFab(ba, dm_, 1, 0);

    res_ = MultiFab(ba, dm_, 1, 0);

  }


  MultiFab& rhs() { return rhs_; }

  MultiFab& phi() { return phi_; }

  const Geometry& geom() const { return geom_; }


  // Solve lap(phi) = rhs in place: local slab -> flat array -> PoissonFFT (internal MPI transpose)

  // -> rewrite the local slab. Mode k=0 set to zero (phi with zero mean).


  void solve() {

    const ConstArray4 f = rhs_.fab(0).const_array();

    const Box2D v = rhs_.box(0);  // local slab [0..Nx-1] x [y0..y0+nyl-1]

    std::vector<double> rho(static_cast<std::size_t>(nyl_) * Nx_), phil;

    for (int jl = 0; jl < nyl_; ++jl)

      for (int i = 0; i < Nx_; ++i)

        rho[static_cast<std::size_t>(jl) * Nx_ + i] = f(v.lo[0] + i, v.lo[1] + jl);

    fft_.solve(rho, phil);

    Array4 p = phi_.fab(0).array();

    for (int jl = 0; jl < nyl_; ++jl)

      for (int i = 0; i < Nx_; ++i)

        p(v.lo[0] + i, v.lo[1] + jl) = phil[static_cast<std::size_t>(jl) * Nx_ + i];

  }


  // Discrete residual ||lap(phi) - rhs|| reduced over all ranks (~round-off: direct solve exact).


  Real residual() {

    BCRec bc;                                                    // periodic

    fill_boundary(phi_, geom_.domain, Periodicity{true, true});  // inter-slab halos (MPI)

    poisson_residual(phi_, rhs_, geom_, bc, res_);

    return all_reduce_max(norm_inf(res_));

  }


 private:

  Geometry geom_;

  int Nx_, nyl_;

  DistributionMapping dm_;

  MultiFab phi_, rhs_, res_;

  PoissonFFT fft_;

};


static_assert(EllipticSolver<DistributedFFTSolver>,

              "DistributedFFTSolver must model EllipticSolver");


class RemappedFFTSolver {

 public:


  RemappedFFTSolver(const Geometry& geom, const BoxArray& ba, const BCRec& = BCRec{},

                    std::function<bool(Real, Real)> = {}, bool spectral = false)

      : geom_(geom),

        dm_(ba.size(), n_ranks()),  // SAME layout System uses: box i -> rank i % np

        phi_(ba, dm_, 1, 1),

        rhs_(ba, dm_, 1, 0),

        res_(ba, dm_, 1, 0),

        Nx_(geom.domain.nx()),

        Ny_(geom.domain.ny()),

        nyl_(geom.domain.ny() / n_ranks()),

        owner_rank_(ba.size() > 0 ? dm_[0] : 0),

        fft_(geom.domain.nx(), geom.domain.ny(), geom.xhi - geom.xlo, geom.yhi - geom.ylo,

             spectral) {

    // CONTRACT (ADC-273 vote on ADC-287): this solver is valid ONLY for System's round-robin

    // single-box Cartesian layout, where box 0 is owned by rank dm_[0] (= owner_rank_) and every other

    // rank holds an empty fab. It presents that layout outward (rhs()/phi() on ba/dm) and hides the

    // box<->slab transpose inside solve(). For a genuinely slab-distributed domain use

    // DistributedFFTSolver; for wall/non-constant-coefficient cases use geometric_mg.

    // COLLECTIVE hard guards (throw on ALL ranks -> no deadlock; ensure_elliptic is collective and

    // every rank constructs this object). Ny must split into np equal slabs for the internal transpose.

    if (ba.size() != 1) {

      throw std::runtime_error(

          "RemappedFFTSolver: System single-box layout required (ba.size()==1); for a "

          "slab-distributed "

          "domain use DistributedFFTSolver or geometric_mg");

    }

    if (geom.domain.ny() % n_ranks() != 0) {

      throw std::runtime_error(

          "RemappedFFTSolver: Ny must be divisible by n_ranks() for the slab FFT");

    }

  }


  MultiFab& rhs() { return rhs_; }

  MultiFab& phi() { return phi_; }

  const Geometry& geom() const { return geom_; }


  // Solve lap(phi) = rhs in place. THREE MPI collectives, in order: (1) MPI_Scatter the owner's box-major

  // rhs into per-rank nyl_ x Nx_ slabs; (2) PoissonFFT::solve runs its internal MPI_Alltoall (y<->x

  // transpose) + FFT on each slab; (3) MPI_Gather the phi slabs back onto the owner. Then the owner fills

  // periodic ghosts on the System layout (intra-box wrap, no messages -> deadlock-free). The caller's

  // device_fence() after ell_solve() (system_field_solver.hpp) covers the GPU read of the centered grad

  // phi, identical to the PoissonFFTSolver path. Mode k=0 set to zero (phi with zero mean).


  void solve() {

    // box-major rhs lives ENTIRELY on owner_rank_; each rank ends up with one nyl_ x Nx_ slab.

    std::vector<double> rho_local(static_cast<std::size_t>(nyl_) * Nx_), phi_local;

#ifdef POPS_HAS_MPI

    std::vector<double> sendbuf;  // only filled on the owner

    if (my_rank() == owner_rank_) {

      const ConstArray4 f = rhs_.fab(0).const_array();

      const Box2D v = rhs_.box(0);  // valid box = whole domain [0..Nx-1] x [0..Ny-1]

      sendbuf.resize(static_cast<std::size_t>(Nx_) * Ny_);

      // pack per destination rank r: rows [r*nyl_, (r+1)*nyl_) contiguous, row-major full x.

      for (int r = 0; r < n_ranks(); ++r)

        for (int jl = 0; jl < nyl_; ++jl)

          for (int i = 0; i < Nx_; ++i)

            sendbuf[(static_cast<std::size_t>(r) * nyl_ + jl) * Nx_ + i] =

                f(v.lo[0] + i, v.lo[1] + r * nyl_ + jl);

    }

    const int cnt = nyl_ * Nx_;

    MPI_Scatter(my_rank() == owner_rank_ ? sendbuf.data() : nullptr, cnt, MPI_DOUBLE,

                rho_local.data(), cnt, MPI_DOUBLE, owner_rank_, MPI_COMM_WORLD);

#else

    {

      const ConstArray4 f = rhs_.fab(0).const_array();

      const Box2D v = rhs_.box(0);

      for (int j = 0; j < Ny_; ++j)

        for (int i = 0; i < Nx_; ++i)

          rho_local[static_cast<std::size_t>(j) * Nx_ + i] = f(v.lo[0] + i, v.lo[1] + j);

    }

#endif

    fft_.solve(rho_local, phi_local);  // PROVEN slab FFT (internal MPI_Alltoall), zero-mean phi

#ifdef POPS_HAS_MPI

    std::vector<double> recvbuf;

    if (my_rank() == owner_rank_)

      recvbuf.resize(static_cast<std::size_t>(Nx_) * Ny_);

    MPI_Gather(phi_local.data(), cnt, MPI_DOUBLE,

               my_rank() == owner_rank_ ? recvbuf.data() : nullptr, cnt, MPI_DOUBLE, owner_rank_,

               MPI_COMM_WORLD);

    if (my_rank() == owner_rank_) {

      Array4 p = phi_.fab(0).array();

      const Box2D v = phi_.box(0);

      for (int r = 0; r < n_ranks(); ++r)

        for (int jl = 0; jl < nyl_; ++jl)

          for (int i = 0; i < Nx_; ++i)

            p(v.lo[0] + i, v.lo[1] + r * nyl_ + jl) =

                recvbuf[(static_cast<std::size_t>(r) * nyl_ + jl) * Nx_ + i];

    }

#else

    {

      Array4 p = phi_.fab(0).array();

      const Box2D v = phi_.box(0);

      for (int j = 0; j < Ny_; ++j)

        for (int i = 0; i < Nx_; ++i)

          p(v.lo[0] + i, v.lo[1] + j) = phi_local[static_cast<std::size_t>(j) * Nx_ + i];

    }

#endif

    // Periodic ghosts on the System layout: only the owner has a fab -> pure intra-box wrap (self-copy),

    // no inter-rank messages, deadlock-free. SAME role as in PoissonFFTSolver::solve() (the centered grad

    // phi of the aux derivation reads the i+-1 / j+-1 ghosts).

    fill_boundary(phi_, geom_.domain, Periodicity{true, true});

    // PR #254 managed-buffer / device-ordering discipline: owner-only device_fence so the host gather +

    // ghost wrap are settled before phi() is read on the device. The caller's device_fence() after

    // ell_solve() (system_field_solver.hpp) already brackets the grad-phi read, so this is belt-and-

    // suspenders, but it self-documents the ordering at the solver seam. Gated on the owner because only

    // it holds a fab (no-op on the empty ranks); no-op on the CPU/Serial backend.

    if (my_rank() == owner_rank_)

      device_fence();

  }


  // Discrete residual ||lap(phi) - rhs|| reduced over all ranks (~round-off: direct solve exact). The

  // box is owner-only, so poisson_residual runs on the owner's fab and is 0 elsewhere; the all_reduce_max

  // is COLLECTIVE (every rank participates).


  Real residual() {

    BCRec bc;  // periodic

    poisson_residual(phi_, rhs_, geom_, bc, res_);

    return all_reduce_max(norm_inf(res_));

  }


 private:

  Geometry geom_;

  DistributionMapping dm_;

  MultiFab phi_, rhs_, res_;

  int Nx_, Ny_, nyl_, owner_rank_;

  PoissonFFT fft_;

};


static_assert(EllipticSolver<RemappedFFTSolver>, "RemappedFFTSolver must model EllipticSolver");


}  // namespace pops

box_array.hpp
BoxArray: the set of boxes tiling a level (disjoint, covering).

pops::BoxArray
Ordered list of boxes tiling a level.
Definition box_array.hpp:22

pops::BoxArray::size
int size() const
Number of boxes in the tiling.
Definition box_array.hpp:42

pops::DistributedFFTSolver
DIRECT periodic Poisson solver (spectral FFT) DISTRIBUTED, models EllipticSolver.
Definition poisson_fft_solver.hpp:143

pops::DistributedFFTSolver::phi
MultiFab & phi()
Definition poisson_fft_solver.hpp:168

pops::DistributedFFTSolver::solve
void solve()
Definition poisson_fft_solver.hpp:173

pops::DistributedFFTSolver::DistributedFFTSolver
DistributedFFTSolver(const Geometry &geom, const BCRec &=BCRec{}, std::function< bool(Real, Real)>={})
Definition poisson_fft_solver.hpp:145

pops::DistributedFFTSolver::residual
Real residual()
Definition poisson_fft_solver.hpp:188

pops::DistributedFFTSolver::rhs
MultiFab & rhs()
Definition poisson_fft_solver.hpp:167

pops::DistributedFFTSolver::geom
const Geometry & geom() const
Definition poisson_fft_solver.hpp:169

pops::DistributionMapping
Owning MPI rank of each box, indexed by GLOBAL box index (parallel to a BoxArray).
Definition distribution_mapping.hpp:19

pops::Fab2D::const_array
ConstArray4 const_array() const
READ handle (POD device-copyable) over this Fab. Valid as long as the Fab lives.
Definition fab2d.hpp:96

pops::Fab2D::array
Array4 array()
WRITE handle (POD device-copyable) over this Fab. Valid as long as the Fab lives.
Definition fab2d.hpp:91

pops::MultiFab
Field distributed over a level: decomposition (BoxArray) + distribution (DistributionMapping) + ncomp...
Definition multifab.hpp:33

pops::MultiFab::fab
Fab2D & fab(int li)
Local fab at index li (0 <= li < local_size()), for writing.
Definition multifab.hpp:67

pops::MultiFab::box
const Box2D & box(int li) const
VALID box of local fab li.
Definition multifab.hpp:71

pops::PoissonFFTSolver
Definition poisson_fft_solver.hpp:51

pops::PoissonFFTSolver::residual
Real residual()
Definition poisson_fft_solver.hpp:107

pops::PoissonFFTSolver::geom
const Geometry & geom() const
Definition poisson_fft_solver.hpp:82

pops::PoissonFFTSolver::phi
MultiFab & phi()
Definition poisson_fft_solver.hpp:81

pops::PoissonFFTSolver::PoissonFFTSolver
PoissonFFTSolver(const Geometry &geom, const BoxArray &ba, const BCRec &=BCRec{}, std::function< bool(Real, Real)>={}, bool spectral=false)
spectral: Laplacian symbol (false = discrete 5-point stencil, bit-identical default; true = continuou...
Definition poisson_fft_solver.hpp:55

pops::PoissonFFTSolver::rhs
MultiFab & rhs()
Definition poisson_fft_solver.hpp:80

pops::PoissonFFTSolver::solve
void solve()
Definition poisson_fft_solver.hpp:85

pops::PoissonFFT
Definition poisson_fft.hpp:101

pops::PoissonFFT::solve
void solve(const std::vector< double > &rho_local, std::vector< double > &phi_local)
Definition poisson_fft.hpp:126

pops::RemappedFFTSolver
DIRECT periodic Poisson solver (spectral FFT) under MPI, presenting the SYSTEM LAYOUT,...
Definition poisson_fft_solver.hpp:241

pops::RemappedFFTSolver::solve
void solve()
Definition poisson_fft_solver.hpp:288

pops::RemappedFFTSolver::RemappedFFTSolver
RemappedFFTSolver(const Geometry &geom, const BoxArray &ba, const BCRec &=BCRec{}, std::function< bool(Real, Real)>={}, bool spectral=false)
spectral: Laplacian symbol forwarded to PoissonFFT (false = discrete 5-point stencil,...
Definition poisson_fft_solver.hpp:246

pops::RemappedFFTSolver::phi
MultiFab & phi()
Definition poisson_fft_solver.hpp:279

pops::RemappedFFTSolver::residual
Real residual()
Definition poisson_fft_solver.hpp:358

pops::RemappedFFTSolver::rhs
MultiFab & rhs()
Definition poisson_fft_solver.hpp:278

pops::RemappedFFTSolver::geom
const Geometry & geom() const
Definition poisson_fft_solver.hpp:280

comm.hpp
Parallel seam: minimal MPI abstraction (rank/size + collectives) with serial fallback.

distribution_mapping.hpp
DistributionMapping: maps each box (by global index) to its owning MPI rank.

elliptic_solver.hpp
EllipticSolver concept: common contract for elliptic solvers at the MultiFab level (solve D phi = f),...

fill_boundary.hpp
fill_boundary: INTRA-level halo exchange (fills ghosts from neighbors).

geometry.hpp
Geometry: index-space (Box2D) <-> Cartesian physical-space mapping; PolarGeometry: SIBLING for a glob...

mf_arith.hpp
MultiFab arithmetic (saxpy, lincomb, norm_inf, dot) over VALID cells.

multifab.hpp
MultiFab: a field DISTRIBUTED over a level (equivalent of AMReX's MultiFab).

pops
Definition amr_hierarchy.hpp:29

pops::Real
double Real
Definition types.hpp:30

pops::n_ranks
int n_ranks()
Definition comm.hpp:139

pops::device_fence
void device_fence()
Device barrier: waits for in-flight kernels to finish before a HOST access to unified memory.
Definition kokkos_env.hpp:43

pops::fill_boundary
void fill_boundary(MultiFab &mf, const Box2D &domain, Periodicity per={})
BLOCKING halo exchange: begin then end immediately (no overlap).
Definition fill_boundary.hpp:333

pops::my_rank
int my_rank()
Definition comm.hpp:136

pops::all_reduce_max
double all_reduce_max(double x)
Definition comm.hpp:146

pops::poisson_residual
void poisson_residual(MultiFab &phi, const MultiFab &f, const Geometry &geom, const BCRec &bc, MultiFab &res, const MultiFab *mask=nullptr, const MultiFab *coef=nullptr, const MultiFab *eps=nullptr, const MultiFab *kappa=nullptr, const MultiFab *eps_y=nullptr, const MultiFab *a_xy=nullptr, const MultiFab *a_yx=nullptr)
Definition poisson_operator.hpp:271

pops::norm_inf
Real norm_inf(const MultiFab &mf, int comp=0)
Infinity norm max |f(.,.,comp)| over the valid cells (LOCAL, without MPI all_reduce).
Definition mf_arith.hpp:123

physical_bc.hpp
PHYSICAL boundary conditions at the domain edge (BCType, BCRec, fill_physical_bc, fill_ghosts).

poisson_fft.hpp
Low-level FFT brick: PoissonFFT (spectral periodic Poisson solver, slab-distributed) plus 1D radix-2 ...

poisson_operator.hpp
Free functions of the elliptic operator: apply_laplacian (matvec), poisson_residual (residual),...

pops::Array4
WRITE POD handle (raw pointer + strides) over a Fab2D buffer, indexed by (i, j, c) IN GLOBAL INDICES ...
Definition fab2d.hpp:29

pops::BCRec
Boundary conditions for the FOUR faces of the domain (type + associated Dirichlet value).
Definition physical_bc.hpp:29

pops::Box2D
2D integer index space, cell-centered.
Definition box2d.hpp:37

pops::Box2D::nx
POPS_HD int nx() const
Width (direction 0). POPS_HD (called from Geometry::dx() in a device kernel).
Definition box2d.hpp:50

pops::Box2D::ny
POPS_HD int ny() const
Height (direction 1). POPS_HD (called from Geometry::dy() in a device kernel).
Definition box2d.hpp:52

pops::Box2D::lo
int lo[2]
Definition box2d.hpp:38

pops::ConstArray4
READ-only handle (const counterpart of Array4): same layout and same contract (POD device-copyable,...
Definition fab2d.hpp:44

pops::Geometry
Cartesian geometry of a level: index domain + physical bounds [xlo, xhi] x [ylo, yhi].
Definition geometry.hpp:20

pops::Geometry::domain
Box2D domain
Definition geometry.hpp:21

pops::Periodicity
Per-direction periodicity: halo wrapping in x and/or y during the exchange (false = open edge,...
Definition fill_boundary.hpp:37