PoPS/doxygen/polar__tensor__operator_8hpp_source.html

#pragma once


#include <pops/core/foundation/types.hpp>

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

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

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

#include <pops/mesh/execution/for_each.hpp>

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

#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 <cmath>

#include <concepts>

#include <cstddef>

#include <stdexcept>

#include <vector>


namespace pops {


namespace detail {


POPS_HD inline Real polar_radial_div(const ConstArray4& p, const ConstArray4& arr, bool hrt,

                                    const ConstArray4& art, int i, int j, Real ri, Real rfm,

                                    Real rfp, Real idr, Real idth) {

  // Face a_rr coefficients (arithmetic mean of the adjacent centers).

  const Real arr_p = Real(0.5) * (arr(i, j) + arr(i + 1, j));

  const Real arr_m = Real(0.5) * (arr(i, j) + arr(i - 1, j));

  // DIAGONAL radial term: (1/r_i) [ r_{i+1/2} a_rr_p (phi_{i+1}-phi_i)/dr - r_{i-1/2} a_rr_m (phi_i-phi_{i-1})/dr ] / dr.

  Real out = (rfp * arr_p * (p(i + 1, j) - p(i, j)) - rfm * arr_m * (p(i, j) - p(i - 1, j))) *

             (idr * idr / ri);

  if (hrt) {  // CROSS radial term a_rt (1/r_face) d_theta phi at face i+-1/2 (d_theta averaged over 4 corners).

    const Real art_p = Real(0.5) * (art(i, j) + art(i + 1, j));

    const Real art_m = Real(0.5) * (art(i, j) + art(i - 1, j));

    // (d_theta phi)_face averaged over the 4 corners of the face: 1/(r_face) because grad_theta = (1/r) d_theta phi.

    const Real dth_p =

        (p(i, j + 1) + p(i + 1, j + 1) - p(i, j - 1) - p(i + 1, j - 1)) * (Real(0.25) * idth);

    const Real dth_m =

        (p(i - 1, j + 1) + p(i, j + 1) - p(i - 1, j - 1) - p(i, j - 1)) * (Real(0.25) * idth);

    // (1/r_i)(1/dr) [ r_{i+1/2} F_cross_p - r_{i-1/2} F_cross_m ], F_cross = a_rt (1/r_face) d_theta phi.

    // r_face * (1/r_face) = 1: the r_face metric cancels the 1/r_face of the azimuthal gradient -> factor 1.

    out += (art_p * dth_p - art_m * dth_m) * (idr / ri);

  }

  return out;

}


POPS_HD inline Real polar_azimuthal_div(const ConstArray4& p, const ConstArray4& att, bool htr,

                                       const ConstArray4& atr, int i, int j, Real ri, Real idr,

                                       Real idth) {

  const Real inv_r2 = Real(1) / (ri * ri);

  const Real att_p = Real(0.5) * (att(i, j) + att(i, j + 1));

  const Real att_m = Real(0.5) * (att(i, j) + att(i, j - 1));

  // DIAGONAL azimuthal term: a_tt^{j+-1/2} (phi_{j+1}-2 phi+phi_{j-1}) / (r_i^2 dtheta^2), flux form.

  Real out =

      (att_p * (p(i, j + 1) - p(i, j)) - att_m * (p(i, j) - p(i, j - 1))) * (idth * idth * inv_r2);

  if (htr) {  // CROSS azimuthal term a_tr d_r phi at face j+-1/2 (d_r averaged over 4 corners).

    const Real atr_p = Real(0.5) * (atr(i, j) + atr(i, j + 1));

    const Real atr_m = Real(0.5) * (atr(i, j) + atr(i, j - 1));

    const Real dr_p =

        (p(i + 1, j) + p(i + 1, j + 1) - p(i - 1, j) - p(i - 1, j + 1)) * (Real(0.25) * idr);

    const Real dr_m =

        (p(i + 1, j - 1) + p(i + 1, j) - p(i - 1, j - 1) - p(i - 1, j)) * (Real(0.25) * idr);

    // (1/r_i)(1/dtheta) [ F_cross_p - F_cross_m ], F_cross = a_tr d_r phi. Metric 1/r_i (face weight dr).

    out += (atr_p * dr_p - atr_m * dr_m) * (idth / ri);

  }

  return out;

}


POPS_HD inline Real polar_diag(const ConstArray4& arr, const ConstArray4& att, int i, int j, Real ri,

                              Real rfm, Real rfp, Real idr, Real idth) {

  const Real arr_p = Real(0.5) * (arr(i, j) + arr(i + 1, j));

  const Real arr_m = Real(0.5) * (arr(i, j) + arr(i - 1, j));

  const Real att_p = Real(0.5) * (att(i, j) + att(i, j + 1));

  const Real att_m = Real(0.5) * (att(i, j) + att(i, j - 1));

  const Real inv_r2 = Real(1) / (ri * ri);

  const Real rad = (rfp * arr_p + rfm * arr_m) * (idr * idr / ri);

  const Real azi = (att_p + att_m) * (idth * idth * inv_r2);

  return -(rad + azi);  // diagonal coefficient of L_int (the neighbors have positive coefficients)

}


struct PolarApplyKernel {

  ConstArray4 p;

  Array4 L;

  ConstArray4 arr, att;

  bool hrt, htr;

  ConstArray4 art, atr;

  Real r_min, dr, idr, idth;


  POPS_HD void operator()(int i, int j) const {

    const Real ri = r_min + (i + Real(0.5)) * dr;  // r_cell(i)

    const Real rfm = r_min + i * dr;               // r_face(i)   = r_{i-1/2}

    const Real rfp = r_min + (i + 1) * dr;         // r_face(i+1) = r_{i+1/2}

    L(i, j) = polar_radial_div(p, arr, hrt, art, i, j, ri, rfm, rfp, idr, idth) +

              polar_azimuthal_div(p, att, htr, atr, i, j, ri, idr, idth);

  }


};


struct PolarJacobiApplyKernel {

  ConstArray4 z;

  Array4 out;

  ConstArray4 idiag;  // 1 / diagonal coefficient of L_int (negative), precomputed

  POPS_HD void operator()(int i, int j) const { out(i, j) = z(i, j) * idiag(i, j); }

};


struct PolarCopyKernel {

  Array4 d;

  ConstArray4 s;

  POPS_HD void operator()(int i, int j) const { d(i, j) = s(i, j, 0); }

};


struct PolarInvDiagKernel {

  ConstArray4 arr, att;

  Array4 idiag;

  Real r_min, dr, idr, idth;


  POPS_HD void operator()(int i, int j) const {

    const Real ri = r_min + (i + Real(0.5)) * dr;

    const Real rfm = r_min + i * dr;

    const Real rfp = r_min + (i + 1) * dr;

    const Real d = polar_diag(arr, att, i, j, ri, rfm, rfp, idr, idth);

    idiag(i, j) = d != Real(0) ? Real(1) / d : Real(0);

  }


};


}  // namespace detail


inline void apply_polar_tensor(const MultiFab& phi, const PolarGeometry& geom, MultiFab& lap,

                               const MultiFab* a_rr, const MultiFab* a_tt, const MultiFab* a_rt,

                               const MultiFab* a_tr) {

  // CONTRACT: a_rr/a_tt required (the diagonal coefficients of the stencil are always read). An

  // isotropic case provides fields CONSTANT at 1 (PolarTensorKrylovSolver does this via its internal

  // stores). a_rt/a_tr optional (cross terms). We cannot deref a nullptr in the kernel -> guard at entry.

  if (!a_rr || !a_tt)

    throw std::runtime_error(

        "apply_polar_tensor: a_rr and a_tt required (fields at 1 if isotropic)");

  const Real dr = geom.dr();

  const Real idr = Real(1) / dr;

  const Real idth = Real(1) / geom.dtheta();

  for (int li = 0; li < phi.local_size(); ++li) {

    const ConstArray4 p = phi.fab(li).const_array();

    Array4 L = lap.fab(li).array();

    const Box2D v = lap.box(li);

    const ConstArray4 arr = a_rr->fab(li).const_array();

    const ConstArray4 att = a_tt->fab(li).const_array();

    const bool hrt = a_rt != nullptr;

    const bool htr = a_tr != nullptr;

    const ConstArray4 art = hrt ? a_rt->fab(li).const_array() : ConstArray4{};

    const ConstArray4 atr = htr ? a_tr->fab(li).const_array() : ConstArray4{};

    for_each_cell(

        v, detail::PolarApplyKernel{p, L, arr, att, hrt, htr, art, atr, geom.r_min, dr, idr, idth});

  }

}


template <class S>


concept PolarLinearSolver = requires(S s, Real tol, int it) {

  { s.rhs() } -> std::same_as<MultiFab&>;

  { s.phi() } -> std::same_as<MultiFab&>;

  s.solve();

  { s.residual() } -> std::convertible_to<Real>;

  { s.geom() } -> std::convertible_to<const PolarGeometry&>;

  s.solve(tol, it);

  requires !std::same_as<decltype(s.solve(tol, it)), void>;

};


struct PolarKrylovResult {

  int iters = 0;

  Real rel_residual = 0;

  bool converged = false;

};


enum class PolarPrecond { Jacobi, RadialLine };


class PolarTensorKrylovSolver {

 public:


  PolarTensorKrylovSolver(const PolarGeometry& geom, const BoxArray& ba, const BCRec& bc,

                          PolarPrecond precond = PolarPrecond::RadialLine)

      : geom_(geom),

        bc_(force_theta_periodic(bc)),

        precond_(precond),

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

        phi_(ba, dm_, 1, 1),

        rhs_(ba, dm_, 1, 0),

        r_(ba, dm_, 1, 0),

        rhat_(ba, dm_, 1, 0),

        p_(ba, dm_, 1, 0),

        v_(ba, dm_, 1, 0),

        s_(ba, dm_, 1, 0),

        t_(ba, dm_, 1, 0),

        phat_(ba, dm_, 1, 1),

        shat_(ba, dm_, 1, 1),

        idiag_(ba, dm_, 1, 0),

        op_offset_(ba, dm_, 1, 0),

        a_rr_store_(ba, dm_, 1, 1),

        a_tt_store_(ba, dm_, 1, 1) {

    // LAYOUT GUARD (replaces the single-rank guard): the RadialLine preconditioner solves a RADIAL

    // tridiagonal per theta line via a SEQUENTIAL Thomas sweep in r. This sweep must stay LOCAL to a

    // box (it cannot cross a box/rank boundary). We therefore REQUIRE that each box of the BoxArray

    // cover the FULL RADIAL RANGE [domain.lo[0], domain.hi[0]] (splitting in theta only). The Jacobi

    // fallback (per cell) has no such constraint -> no check. Single-rank single box: the box covers

    // the whole ring, the check passes trivially (path unchanged).

    if (precond_ == PolarPrecond::RadialLine)

      check_radial_columns(ba);

    // Default diagonal coefficients = 1 (isotropic): internal fields at 1, ghosts filled Foextrap.

    a_rr_store_.set_val(Real(1));

    a_tt_store_.set_val(Real(1));

    a_rr_ = &a_rr_store_;

    a_tt_ = &a_tt_store_;

  }


  // RULE OF FIVE (C.21): the current pointers a_rr_/a_tt_ alias the internal stores

  // a_rr_store_/a_tt_store_ or external fields. A DEFAULT copy/move would leave these pointers aiming

  // at the SOURCE object's stores (dangling/UB). The solver is ALWAYS used as a LOCAL scope variable

  // (never copied, moved, stored in a container, nor returned by value): we DELETE the four

  // operations rather than write a move re-pointing the stores (useless here).

  PolarTensorKrylovSolver(const PolarTensorKrylovSolver&) = delete;

  PolarTensorKrylovSolver& operator=(const PolarTensorKrylovSolver&) = delete;

  PolarTensorKrylovSolver(PolarTensorKrylovSolver&&) = delete;

  PolarTensorKrylovSolver& operator=(PolarTensorKrylovSolver&&) = delete;


  // --- PolarEllipticSolver / PolarLinearSolver contract ---

  MultiFab& rhs() { return rhs_; }

  MultiFab& phi() { return phi_; }

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


  void set_coefficients(MultiFab* a_rr, MultiFab* a_tt, MultiFab* a_rt = nullptr,

                        MultiFab* a_tr = nullptr) {

    a_rr_ = a_rr ? a_rr : &a_rr_store_;

    a_tt_ = a_tt ? a_tt : &a_tt_store_;

    a_rt_ = a_rt;

    a_tr_ = a_tr;

    fill_coeff_ghosts();

    coeffs_ready_ = false;  // idiag to be recomputed

  }


  Real residual() {

    ensure_coeffs();

    apply_operator(phi_, r_);

    lincomb(r_, Real(1), rhs_, Real(-1), r_);

    return l2_norm(r_);

  }


  void solve() { solve(Real(1e-10), 400); }


  PolarKrylovResult solve(Real rel_tol, int max_iters) {

    ensure_coeffs();

    prepare_offset();  // c_bc = apply_operator(0) (inhomogeneous part of Dirichlet boundary) once

    PolarKrylovResult res;

    // MULTI-RANK MPI: ALL ranks (including those without a box, local_size()==0) execute the SAME

    // BiCGStab body. The per-fab operations (lincomb/saxpy/copy/apply_precond/apply_polar_tensor) are

    // no-ops on an empty rank; the COLLECTIVES (dot/l2_norm/project_mean -> all_reduce_sum) are called

    // by ALL ranks in the SAME order (an empty rank contributes 0). The stopping criteria rely only on

    // these GLOBAL scalars (identical everywhere) -> no desynchronization, no separate "run_empty_rank"

    // path to maintain in parallel. fill_ghosts (fill_boundary) is a PAIRWISE exchange (paired

    // Isend/Irecv, not a collective): an empty rank posts nothing and nobody sends to it -> no deadlock.

    // r0 = rhs - L_int(phi) (AFFINE operator: the Dirichlet data is folded into the residual).

    apply_operator(phi_, v_);

    lincomb(r_, Real(1), rhs_, Real(-1), v_);

    // SINGULAR OPERATOR (pure radial Neumann + periodic theta, no reaction): the constant is in the

    // kernel of L_int. BiCGStab then diverges (the residual keeps an undamped constant component). We

    // FIX THE GAUGE by projection onto the zero-mean subspace: we remove the mean of r0 (the RHS must

    // be compatible: zero mean in the FV sense) and of phi (gauge), then of each preconditioned

    // correction direction in the loop. This is the iterative counterpart of the mode-0 pinning of the

    // direct PolarPoissonSolver, without perturbing the stencil. Dirichlet case (>= one boundary):

    // pin_gauge_ stays false -> PATH UNCHANGED, bit-identical.

    if (pin_gauge_) {

      project_mean(r_);

      project_mean(phi_);

    }

    const Real bnorm = l2_norm(rhs_);

    const Real norm0 = bnorm > Real(0) ? bnorm : Real(1);

    Real rnorm = l2_norm(r_);

    res.rel_residual = rnorm / norm0;

    if (rnorm <= rel_tol * norm0) {

      res.converged = true;

      return res;

    }


    copy_into(rhat_, r_);

    p_.set_val(Real(0));

    v_.set_val(Real(0));

    Real rho_prev = Real(1), alpha = Real(1), omega = Real(1);


    for (int k = 1; k <= max_iters; ++k) {

      const Real rho = dot(rhat_, r_);  // COLLECTIVE

      if (std::fabs(rho) < kTiny || std::fabs(omega) < kTiny) {

        res.iters = k - 1;

        res.rel_residual = rnorm / norm0;

        return res;

      }

      const Real beta = (rho / rho_prev) * (alpha / omega);

      lincomb(p_, Real(1), p_, -omega, v_);  // p <- p - omega v

      lincomb(p_, beta, p_, Real(1), r_);    // p <- r + beta p

      apply_precond(p_, phat_);              // phat = M^{-1} p

      if (pin_gauge_)

        project_mean(phat_);          // gauge: zero-mean correction direction

      apply_operator_lin(phat_, v_);  // v = L_lin(phat) (LINEAR matvec)

      const Real rhat_dot_v = dot(rhat_, v_);

      if (std::fabs(rhat_dot_v) < kTiny) {

        res.iters = k - 1;

        res.rel_residual = rnorm / norm0;

        return res;

      }

      alpha = rho / rhat_dot_v;

      lincomb(s_, Real(1), r_, -alpha, v_);  // s <- r - alpha v

      saxpy(phi_, alpha, phat_);             // phi <- phi + alpha phat

      const Real snorm = l2_norm(s_);

      if (snorm <= rel_tol * norm0) {

        rnorm = snorm;

        res.iters = k;

        res.rel_residual = rnorm / norm0;

        res.converged = true;

        return res;

      }

      apply_precond(s_, shat_);  // shat = M^{-1} s

      if (pin_gauge_)

        project_mean(shat_);          // gauge: zero-mean correction direction

      apply_operator_lin(shat_, t_);  // t = L_lin(shat)

      const Real tt = dot(t_, t_);

      omega = tt > kTiny ? dot(t_, s_) / tt : Real(0);

      saxpy(phi_, omega, shat_);             // phi <- phi + omega shat

      lincomb(r_, Real(1), s_, -omega, t_);  // r <- s - omega t

      rnorm = l2_norm(r_);

      res.iters = k;

      res.rel_residual = rnorm / norm0;

      if (rnorm <= rel_tol * norm0) {

        res.converged = true;

        return res;

      }

      rho_prev = rho;

    }

    return res;  // max_iters reached: best effort (converged=false)

  }


 private:

  static constexpr Real kTiny = Real(1e-300);


  void check_radial_columns(const BoxArray& ba) const {

    const Box2D dom = geom_.domain;

    for (int b = 0; b < ba.size(); ++b) {

      if (ba[b].lo[0] != dom.lo[0] || ba[b].hi[0] != dom.hi[0])

        throw std::runtime_error(

            "PolarTensorKrylovSolver (precond RadialLine): the BoxArray must be split in THETA "

            "only "

            "(each box covers the full radial range). The Thomas sweep in r cannot cross a box "

            "boundary. Use an azimuthal splitting, or the PolarPrecond::Jacobi fallback (per cell, "

            "without layout constraint) if r must be cut.");

    }

  }


  static BCRec force_theta_periodic(const BCRec& bc) {

    BCRec b = bc;

    b.ylo = BCType::Periodic;

    b.yhi = BCType::Periodic;

    b.ylo_val = Real(0);

    b.yhi_val = Real(0);

    return b;

  }


  BCRec coeff_bc() const {

    auto fo = [](BCType t) { return t == BCType::Periodic ? t : BCType::Foextrap; };

    BCRec b;

    b.xlo = fo(bc_.xlo);

    b.xhi = fo(bc_.xhi);

    b.ylo = BCType::Periodic;

    b.yhi = BCType::Periodic;  // theta always periodic

    return b;

  }


  void fill_coeff_ghosts() {

    const BCRec eb = coeff_bc();

    device_fence();

    fill_ghosts(*a_rr_, geom_.domain, eb);

    fill_ghosts(*a_tt_, geom_.domain, eb);

    if (a_rt_)

      fill_ghosts(*a_rt_, geom_.domain, eb);

    if (a_tr_)

      fill_ghosts(*a_tr_, geom_.domain, eb);

  }


  void ensure_coeffs() {

    if (coeffs_ready_)

      return;

    // idiag = 1/diag of the diagonal stencil (for Jacobi).

    const Real dr = geom_.dr();

    const Real idr = Real(1) / dr;

    const Real idth = Real(1) / geom_.dtheta();

    for (int li = 0; li < idiag_.local_size(); ++li) {

      for_each_cell(idiag_.box(li), detail::PolarInvDiagKernel{

                                        a_rr_->fab(li).const_array(), a_tt_->fab(li).const_array(),

                                        idiag_.fab(li).array(), geom_.r_min, dr, idr, idth});

    }

    if (precond_ == PolarPrecond::RadialLine)

      build_radial_lines();

    coeffs_ready_ = true;

  }


  void build_radial_lines() {

    a_rr_->sync_host();

    a_tt_->sync_host();

    const int nr = geom_.domain.nx();

    const Box2D dom = geom_.domain;

    const Real dr = geom_.dr();

    const Real idr2 = Real(1) / (dr * dr);

    const Real idth2 = Real(1) / (geom_.dtheta() * geom_.dtheta());

    const bool dir_lo = bc_.xlo == BCType::Dirichlet;

    const bool dir_hi = bc_.xhi == BCType::Dirichlet;

    nr_ = nr;

    const int nloc = idiag_.local_size();

    line_b_.resize(static_cast<std::size_t>(nloc));

    line_sub_.resize(static_cast<std::size_t>(nloc));

    line_sup_.resize(static_cast<std::size_t>(nloc));

    for (int li = 0; li < nloc; ++li) {

      const Box2D vb = idiag_.box(li);  // local VALID box (full r range, theta sub-range)

      const int nth_l = vb.ny();        // number of theta lines LOCAL to this box

      const ConstArray4 arr = a_rr_->fab(li).const_array();

      const ConstArray4 att = a_tt_->fab(li).const_array();

      std::vector<Real>& lb = line_b_[li];

      std::vector<Real>& lsub = line_sub_[li];

      std::vector<Real>& lsup = line_sup_[li];

      lb.assign(static_cast<std::size_t>(nth_l) * static_cast<std::size_t>(nr), Real(0));

      lsub.assign(static_cast<std::size_t>(nth_l) * static_cast<std::size_t>(nr), Real(0));

      lsup.assign(static_cast<std::size_t>(nth_l) * static_cast<std::size_t>(nr), Real(0));

      for (int jl = 0; jl < nth_l; ++jl) {

        const int jg = vb.lo[1] + jl;  // GLOBAL theta index

        for (int i = 0; i < nr; ++i) {

          const int ig = dom.lo[0] + i;

          const Real ri = geom_.r_cell(i);

          const Real rfm = geom_.r_face(i);

          const Real rfp = geom_.r_face(i + 1);

          const Real arr_p = Real(0.5) * (arr(ig, jg) + arr(ig + 1, jg));

          const Real arr_m = Real(0.5) * (arr(ig, jg) + arr(ig - 1, jg));

          const Real att_p = Real(0.5) * (att(ig, jg) + att(ig, jg + 1));

          const Real att_m = Real(0.5) * (att(ig, jg) + att(ig, jg - 1));

          const Real ai = rfm * arr_m * (idr2 / ri);  // coeff of p_{i-1} in L_int

          const Real ci = rfp * arr_p * (idr2 / ri);  // coeff of p_{i+1}

          const Real azi_diag =

              (att_p + att_m) * (idth2 / (ri * ri));  // azimuthal part of -diag (lumped)

          Real bi = -(ai + ci) - azi_diag;            // FULL diagonal of L_int (radial + azimuthal)

          Real sub = ai, sup = ci;

          if (i ==

              0) {  // HOMOGENEOUS low BC fold: Dirichlet b -= a, Neumann b += a; sub-diag zeroed

            bi += dir_lo ? -ai : ai;

            sub = Real(0);

          }

          if (i == nr - 1) {  // HOMOGENEOUS high BC fold

            bi += dir_hi ? -ci : ci;

            sup = Real(0);

          }

          const std::size_t idx = static_cast<std::size_t>(jl) * static_cast<std::size_t>(nr) +

                                  static_cast<std::size_t>(i);

          lb[idx] = bi;

          lsub[idx] = sub;

          lsup[idx] = sup;

        }

      }

    }

  }


  void apply_operator(MultiFab& in, MultiFab& out) {

    device_fence();

    fill_ghosts(in, geom_.domain, bc_);

    apply_polar_tensor(in, geom_, out, a_rr_, a_tt_, a_rt_, a_tr_);

  }


  void apply_operator_lin(MultiFab& in, MultiFab& out) {

    apply_operator(in, out);

    if (has_op_offset_)

      lincomb(out, Real(1), out, Real(-1), op_offset_);

  }


  void apply_precond(MultiFab& in, MultiFab& out) {

    if (precond_ == PolarPrecond::Jacobi) {

      for (int li = 0; li < out.local_size(); ++li)

        for_each_cell(out.box(li),

                      detail::PolarJacobiApplyKernel{in.fab(li).const_array(), out.fab(li).array(),

                                                     idiag_.fab(li).const_array()});

      return;

    }

    apply_precond_radial_line(in, out);

  }


  void apply_precond_radial_line(MultiFab& in, MultiFab& out) {

    if (in.local_size() == 0)

      return;

    in.sync_host();

    const int r_lo = geom_.domain.lo[0];

    const std::size_t N = static_cast<std::size_t>(nr_);

    cthom_.assign(N, Real(0));  // Thomas working super-diagonal

    xthom_.assign(N, Real(0));  // working solution

    for (int li = 0; li < in.local_size(); ++li) {

      const Box2D vb = in.box(li);

      const ConstArray4 z = in.fab(li).const_array();

      Array4 o = out.fab(li).array();

      const std::vector<Real>& lb = line_b_[li];

      const std::vector<Real>& lsub = line_sub_[li];

      const std::vector<Real>& lsup = line_sup_[li];

      const int nth_l = vb.ny();

      for (int jl = 0; jl < nth_l; ++jl) {

        const int jg = vb.lo[1] + jl;

        const std::size_t base = static_cast<std::size_t>(jl) * N;

        // Thomas in r: a = lsub, b = lb, c = lsup (base + i), rhs = z(., jg).

        Real beta = lb[base + 0];

        xthom_[0] = (beta != Real(0) ? z(r_lo, jg) / beta : Real(0));

        for (std::size_t i = 1; i < N; ++i) {

          cthom_[i] = lsup[base + i - 1] / beta;

          beta = lb[base + i] - lsub[base + i] * cthom_[i];

          const Real zi = z(r_lo + static_cast<int>(i), jg);

          xthom_[i] = (beta != Real(0) ? (zi - lsub[base + i] * xthom_[i - 1]) / beta : Real(0));

        }

        for (int i = static_cast<int>(N) - 2; i >= 0; --i)

          xthom_[static_cast<std::size_t>(i)] -=

              cthom_[static_cast<std::size_t>(i + 1)] * xthom_[static_cast<std::size_t>(i + 1)];

        for (std::size_t i = 0; i < N; ++i)

          o(r_lo + static_cast<int>(i), jg) = xthom_[i];

      }

    }

  }


  void prepare_offset() {

    has_op_offset_ = (bc_.xlo == BCType::Dirichlet && bc_.xlo_val != Real(0)) ||

                     (bc_.xhi == BCType::Dirichlet && bc_.xhi_val != Real(0));

    pin_gauge_ = (bc_.xlo != BCType::Dirichlet) && (bc_.xhi != BCType::Dirichlet);

    if (has_op_offset_) {

      // Called on ALL ranks (apply_operator is a no-op on an empty rank except for the pairwise

      // fill_ghosts): no local_size() branch that would unpair the cross-rank halo exchange.

      phat_.set_val(Real(0));

      apply_operator(phat_, op_offset_);  // op_offset_ <- L_int(0) = c_bc

    }

  }


  void project_mean(MultiFab& x) {

    x.sync_host();

    const int r_lo = geom_.domain.lo[0];

    const Real dr = geom_.dr(), dth = geom_.dtheta();

    Real sum = 0, vol = 0;

    for (int li = 0; li < x.local_size(); ++li) {

      const Box2D vb = x.box(li);

      const ConstArray4 a = x.fab(li).const_array();

      for (int j = vb.lo[1]; j <= vb.hi[1]; ++j)

        for (int i = vb.lo[0]; i <= vb.hi[0]; ++i) {

          const Real w = geom_.r_cell(i - r_lo) * dr * dth;  // i - r_lo = 0-based radial offset

          sum += a(i, j) * w;

          vol += w;

        }

    }

    // COLLECTIVE all-reduce (over all ranks, including empty ones): GLOBAL mean identical everywhere.

    sum = static_cast<Real>(all_reduce_sum(static_cast<double>(sum)));

    vol = static_cast<Real>(all_reduce_sum(static_cast<double>(vol)));

    const Real mean = vol > Real(0) ? sum / vol : Real(0);

    for (int li = 0; li < x.local_size(); ++li) {

      const Box2D vb = x.box(li);

      Array4 a = x.fab(li).array();

      for (int j = vb.lo[1]; j <= vb.hi[1]; ++j)

        for (int i = vb.lo[0]; i <= vb.hi[0]; ++i)

          a(i, j) -= mean;

    }

  }


  Real l2_norm(const MultiFab& x) { return std::sqrt(dot(x, x)); }


  void copy_into(MultiFab& dst, const MultiFab& src) {

    for (int li = 0; li < dst.local_size(); ++li) {

      Array4 d = dst.fab(li).array();

      const ConstArray4 s = src.fab(li).const_array();

      for_each_cell(dst.box(li), detail::PolarCopyKernel{d, s});

    }

  }


  PolarGeometry geom_;

  BCRec bc_;

  PolarPrecond precond_;

  DistributionMapping dm_;

  MultiFab phi_, rhs_;

  MultiFab r_, rhat_, p_, v_, s_, t_;

  MultiFab phat_, shat_;

  MultiFab idiag_;

  MultiFab op_offset_;

  MultiFab a_rr_store_, a_tt_store_;

  MultiFab* a_rr_ = nullptr;

  MultiFab* a_tt_ = nullptr;

  MultiFab* a_rt_ = nullptr;

  MultiFab* a_tr_ = nullptr;

  bool coeffs_ready_ = false;

  bool has_op_offset_ = false;

  bool pin_gauge_ =

      false;

  // RadialLine preconditioner: radial tridiagonal of the diagonal block per theta line. PER LOCAL BOX

  // (line_b_/sub_/sup_[li]), each stored [jl*nr + i] where jl = theta index LOCAL to the box, i = global

  // radius 0..nr-1 (a_rr may depend on theta -> coeffs per (i, jl)). The splitting being in THETA only,

  // M^{-1} is block-diagonal per box (no cross-box radial coupling) -> each box is inverted

  // independently. cthom_/xthom_ = Thomas working buffers (reused, length nr). HOST.

  std::vector<std::vector<Real>> line_b_, line_sub_,

      line_sup_;

  mutable std::vector<Real> cthom_, xthom_;

  int nr_ = 0;

};


static_assert(PolarLinearSolver<PolarTensorKrylovSolver>,

              "PolarTensorKrylovSolver must model PolarLinearSolver");


}  // 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::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::sync_host
void sync_host()
Makes the HOST residence valid (before a host access: operator(), loop, set_val).
Definition multifab.hpp:79

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

pops::MultiFab::set_val
void set_val(Real v)
Fills all cells (valid + ghosts) of every local fab with v.
Definition multifab.hpp:85

pops::MultiFab::local_size
int local_size() const
Number of fabs OWNED by this rank (bound on local indices).
Definition multifab.hpp:65

pops::PolarTensorKrylovSolver
MATRIX-FREE BiCGStab Krylov solver for the FULL-tensor POLAR elliptic operator L_int(phi) = div(A gra...
Definition polar_tensor_operator.hpp:332

pops::PolarTensorKrylovSolver::operator=
PolarTensorKrylovSolver & operator=(PolarTensorKrylovSolver &&)=delete

pops::PolarTensorKrylovSolver::PolarTensorKrylovSolver
PolarTensorKrylovSolver(const PolarTensorKrylovSolver &)=delete

pops::PolarTensorKrylovSolver::set_coefficients
void set_coefficients(MultiFab *a_rr, MultiFab *a_tt, MultiFab *a_rt=nullptr, MultiFab *a_tr=nullptr)
Sets the tensor A coefficients.
Definition polar_tensor_operator.hpp:391

pops::PolarTensorKrylovSolver::operator=
PolarTensorKrylovSolver & operator=(const PolarTensorKrylovSolver &)=delete

pops::PolarTensorKrylovSolver::rhs
MultiFab & rhs()
Definition polar_tensor_operator.hpp:383

pops::PolarTensorKrylovSolver::residual
Real residual()
Current GLOBAL L2 residual ||rhs - L_int(phi)|| (collective). Prepares once if needed.
Definition polar_tensor_operator.hpp:402

pops::PolarTensorKrylovSolver::geom
const PolarGeometry & geom() const
Definition polar_tensor_operator.hpp:385

pops::PolarTensorKrylovSolver::PolarTensorKrylovSolver
PolarTensorKrylovSolver(PolarTensorKrylovSolver &&)=delete

pops::PolarTensorKrylovSolver::solve
PolarKrylovResult solve(Real rel_tol, int max_iters)
MATRIX-FREE BiCGStab preconditioned by precond_ (RadialLine by default, Jacobi as fallback); fixes th...
Definition polar_tensor_operator.hpp:414

pops::PolarTensorKrylovSolver::phi
MultiFab & phi()
Definition polar_tensor_operator.hpp:384

pops::PolarTensorKrylovSolver::solve
void solve()
Definition polar_tensor_operator.hpp:409

pops::PolarTensorKrylovSolver::PolarTensorKrylovSolver
PolarTensorKrylovSolver(const PolarGeometry &geom, const BoxArray &ba, const BCRec &bc, PolarPrecond precond=PolarPrecond::RadialLine)
Definition polar_tensor_operator.hpp:337

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

pops::PolarLinearSolver
Contract of the iterative POLAR elliptic operators: same shape as PolarEllipticSolver (cf.
Definition polar_tensor_operator.hpp:291

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

fab2d.hpp
Fab2D: single-grid data on a Box2D (in-house equivalent of AMReX's FArrayBox); Array4 / ConstArray4: ...

for_each.hpp
for_each_cell and reductions: the parallelism SEAM over the cells of a Box2D; sync_host / sync_device...

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::detail::polar_diag
POPS_HD Real polar_diag(const ConstArray4 &arr, const ConstArray4 &att, int i, int j, Real ri, Real rfm, Real rfp, Real idr, Real idth)
Diagonal (coefficient of phi_{i,j}) of the diagonal POLAR stencil (radial + azimuthal),...
Definition polar_tensor_operator.hpp:188

pops::detail::polar_radial_div
POPS_HD Real polar_radial_div(const ConstArray4 &p, const ConstArray4 &arr, bool hrt, const ConstArray4 &art, int i, int j, Real ri, Real rfm, Real rfp, Real idr, Real idth)
RADIAL FACE contribution to the diagonal + cross stencil, at (i, j).
Definition polar_tensor_operator.hpp:135

pops::detail::polar_azimuthal_div
POPS_HD Real polar_azimuthal_div(const ConstArray4 &p, const ConstArray4 &att, bool htr, const ConstArray4 &atr, int i, int j, Real ri, Real idr, Real idth)
AZIMUTHAL FACE contribution to the diagonal + cross stencil, at (i, j).
Definition polar_tensor_operator.hpp:162

pops
Definition amr_hierarchy.hpp:29

pops::for_each_cell
void for_each_cell(const Box2D &b, F f)
Applies f to EACH cell (i, j) of box b (bounds inclusive), via Kokkos::parallel_for (Serial / OpenMP ...
Definition for_each.hpp:138

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

pops::dot
Real dot(const MultiFab &x, const MultiFab &y, int comp=0)
Dot product Sum_cells x.y over component comp, reduced over ALL ranks (all_reduce).
Definition mf_arith.hpp:163

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

pops::all_reduce_sum
double all_reduce_sum(double x)
Definition comm.hpp:143

pops::saxpy
void saxpy(MultiFab &y, Real a, const MultiFab &x)
y <- y + a x over ALL components of the valid cells. Identical layouts required.
Definition mf_arith.hpp:102

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::sum
Real sum(const MultiFab &mf, int comp=0)
Sum of the VALID cells of component comp, reduced over ALL ranks (all_reduce).
Definition multifab.hpp:121

pops::apply_polar_tensor
void apply_polar_tensor(const MultiFab &phi, const PolarGeometry &geom, MultiFab &lap, const MultiFab *a_rr, const MultiFab *a_tt, const MultiFab *a_rt, const MultiFab *a_tr)
Applies L_int(phi) = div(A grad phi) in polar over the whole MultiFab.
Definition polar_tensor_operator.hpp:259

pops::lincomb
void lincomb(MultiFab &z, Real a, const MultiFab &x, Real b, const MultiFab &y)
z <- a x + b y over ALL components of the valid cells. Identical layouts; aliasing safe.
Definition mf_arith.hpp:133

pops::PolarPrecond
PolarPrecond
Choice of the SIMPLE BiCGStab PRECONDITIONER (NO MG V-cycle – stagnation on polar 1/r^2,...
Definition polar_tensor_operator.hpp:319

pops::PolarPrecond::Jacobi
@ Jacobi

pops::PolarPrecond::RadialLine
@ RadialLine

pops::BCType
BCType
Boundary condition type for a face: Periodic (handled by fill_boundary), Foextrap (zero gradient,...
Definition physical_bc.hpp:25

pops::BCType::Foextrap
@ Foextrap

pops::BCType::Dirichlet
@ Dirichlet

pops::BCType::Periodic
@ Periodic

pops::fill_ghosts
void fill_ghosts(MultiFab &mf, const Box2D &domain, const BCRec &bc)
COMPLETE ghost filling: fill_boundary (interior + periodic, periodicity deduced from bc) THEN fill_ph...
Definition physical_bc.hpp:227

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

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::BCRec::xlo
BCType xlo
Definition physical_bc.hpp:30

pops::BCRec::xlo_val
Real xlo_val
Definition physical_bc.hpp:32

pops::BCRec::xhi_val
Real xhi_val
Definition physical_bc.hpp:32

pops::BCRec::xhi
BCType xhi
Definition physical_bc.hpp:30

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::hi
int hi[2]
Definition box2d.hpp:39

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::PolarGeometry
Definition geometry.hpp:59

pops::PolarGeometry::r_face
POPS_HD Real r_face(int i) const
Radius at radial FACE i (face between cells i-1 and i; i = 0 -> r_min, i = nr -> r_max).
Definition geometry.hpp:75

pops::PolarGeometry::r_cell
POPS_HD Real r_cell(int i) const
Radius at the CENTER of radial cell i (i = 0 -> r_min + dr/2).
Definition geometry.hpp:73

pops::PolarGeometry::dr
POPS_HD Real dr() const
Definition geometry.hpp:70

pops::PolarGeometry::r_min
Real r_min
Definition geometry.hpp:61

pops::PolarGeometry::domain
Box2D domain
nx() = nr (radial cells), ny() = ntheta (azimuthal cells)
Definition geometry.hpp:60

pops::PolarGeometry::dtheta
POPS_HD Real dtheta() const
Definition geometry.hpp:71

pops::PolarKrylovResult
Result of a polar BiCGStab solve: iterations, relative residual, convergence.
Definition polar_tensor_operator.hpp:304

pops::PolarKrylovResult::converged
bool converged
Definition polar_tensor_operator.hpp:307

pops::PolarKrylovResult::rel_residual
Real rel_residual
Definition polar_tensor_operator.hpp:306

pops::PolarKrylovResult::iters
int iters
Definition polar_tensor_operator.hpp:305

pops::detail::PolarApplyKernel
L_int(phi) = div(A grad phi) in polar (apply).
Definition polar_tensor_operator.hpp:204

pops::detail::PolarApplyKernel::atr
ConstArray4 atr
Definition polar_tensor_operator.hpp:209

pops::detail::PolarApplyKernel::att
ConstArray4 att
Definition polar_tensor_operator.hpp:207

pops::detail::PolarApplyKernel::L
Array4 L
Definition polar_tensor_operator.hpp:206

pops::detail::PolarApplyKernel::dr
Real dr
Definition polar_tensor_operator.hpp:210

pops::detail::PolarApplyKernel::art
ConstArray4 art
Definition polar_tensor_operator.hpp:209

pops::detail::PolarApplyKernel::htr
bool htr
Definition polar_tensor_operator.hpp:208

pops::detail::PolarApplyKernel::operator()
POPS_HD void operator()(int i, int j) const
Definition polar_tensor_operator.hpp:211

pops::detail::PolarApplyKernel::hrt
bool hrt
Definition polar_tensor_operator.hpp:208

pops::detail::PolarApplyKernel::p
ConstArray4 p
Definition polar_tensor_operator.hpp:205

pops::detail::PolarApplyKernel::idth
Real idth
Definition polar_tensor_operator.hpp:210

pops::detail::PolarApplyKernel::idr
Real idr
Definition polar_tensor_operator.hpp:210

pops::detail::PolarApplyKernel::arr
ConstArray4 arr
Definition polar_tensor_operator.hpp:207

pops::detail::PolarApplyKernel::r_min
Real r_min
Definition polar_tensor_operator.hpp:210

pops::detail::PolarCopyKernel
Copy component 0 (dst <- src).
Definition polar_tensor_operator.hpp:234

pops::detail::PolarCopyKernel::operator()
POPS_HD void operator()(int i, int j) const
Definition polar_tensor_operator.hpp:237

pops::detail::PolarCopyKernel::d
Array4 d
Definition polar_tensor_operator.hpp:235

pops::detail::PolarCopyKernel::s
ConstArray4 s
Definition polar_tensor_operator.hpp:236

pops::detail::PolarInvDiagKernel
Computes idiag = 1 / diag of the diagonal polar stencil (for Jacobi). diag = polar_diag (< 0).
Definition polar_tensor_operator.hpp:241

pops::detail::PolarInvDiagKernel::idiag
Array4 idiag
Definition polar_tensor_operator.hpp:243

pops::detail::PolarInvDiagKernel::operator()
POPS_HD void operator()(int i, int j) const
Definition polar_tensor_operator.hpp:245

pops::detail::PolarInvDiagKernel::idr
Real idr
Definition polar_tensor_operator.hpp:244

pops::detail::PolarInvDiagKernel::arr
ConstArray4 arr
Definition polar_tensor_operator.hpp:242

pops::detail::PolarInvDiagKernel::dr
Real dr
Definition polar_tensor_operator.hpp:244

pops::detail::PolarInvDiagKernel::idth
Real idth
Definition polar_tensor_operator.hpp:244

pops::detail::PolarInvDiagKernel::r_min
Real r_min
Definition polar_tensor_operator.hpp:244

pops::detail::PolarInvDiagKernel::att
ConstArray4 att
Definition polar_tensor_operator.hpp:242

pops::detail::PolarJacobiApplyKernel
out = (f - L0 phi) / |diag| – one Jacobi iteration (point-by-point relaxation) on the DIAGONAL polar ...
Definition polar_tensor_operator.hpp:225

pops::detail::PolarJacobiApplyKernel::operator()
POPS_HD void operator()(int i, int j) const
Definition polar_tensor_operator.hpp:229

pops::detail::PolarJacobiApplyKernel::idiag
ConstArray4 idiag
Definition polar_tensor_operator.hpp:228

pops::detail::PolarJacobiApplyKernel::out
Array4 out
Definition polar_tensor_operator.hpp:227

pops::detail::PolarJacobiApplyKernel::z
ConstArray4 z
Definition polar_tensor_operator.hpp:226

types.hpp
Base scalar types and the POPS_HD macro (host+device portability).

POPS_HD
#define POPS_HD
Definition types.hpp:25