flow.hpp

Go to the documentation of this file.

// core — device (Kokkos) collocated incompressible Stokes step on a BlockOctree.
//
// The device counterpart of oracle::AmrFlow (flow.hpp): the whole cut-cell IBM projection step
// runs in Kokkos kernels instead of the host-serial gaussSeidel + host projection the drag
// study found to be the bottleneck. It reuses the host AmrCutCell / AmrPoisson to *build*
// the operators (geometry, openness, cut stencils — done once), then drives the time step
// entirely on the device:
//   * momentum predictor — MomentumOp (assembled cut-cell operator) solved with the
//     parallel BiCGStab of momentum.hpp;
//   * pressure projection — the openness Poisson on Multigrid / PCG;
//   * divergence, ABC gradient correction, rotational pressure update — face-CSR kernels
//     (FaceGeom) that mirror AmrPoisson::forEachFaceFull (same 2:1 sub-faces +
//     openness), so D / G / L stay consistent exactly as in the host collocated coupling.
//
// Stokes (advection off) and Navier–Stokes (setAdvection): implicit-FOU + explicit SOU/Koren
// deferred correction, advected by the divergence-free face field uf (built each projection;
// falls back to ½(u_i+u_j) until the first projection) — conservative (∇·uf = 0).
//
// Requires a Kokkos build + the morton checkout (PECLET_CORE_HAVE_MORTON).
#ifndef PECLET_CORE_AMR_FLOW_HPP
#define PECLET_CORE_AMR_FLOW_HPP
 
#ifdef PECLET_CORE_HAVE_MORTON
 
#include <array>
#include <cmath>
#include <cstdlib>
#include <vector>
 
#include "peclet/core/amr/advect_recon.hpp"  // shared high-order face reconstruction (host+device)
#include "peclet/core/amr/block_octree.hpp"
#include "peclet/core/amr/cut_cell.hpp"
#include "peclet/core/amr/face_geom.hpp"          // FaceGeom (shared with the device assembler)
#include "peclet/core/amr/facegeom_assembly.hpp"  // assembleFaceGeom (D4/D6)
#include "peclet/core/amr/momentum.hpp"
#include "peclet/core/amr/momentum_assembly.hpp"  // assembleMomentum (D3/D6)
#include "peclet/core/amr/multigrid.hpp"
#include "peclet/core/amr/pcg.hpp"
#include "peclet/core/amr/poisson.hpp"
#include "peclet/core/amr/velocity_mg.hpp"
#include "peclet/core/common/types.hpp"
#include "peclet/core/common/view.hpp"
 
namespace peclet::core::amr {
 
// FaceGeom (the collocated projection's static face-geometry CSR) now lives in face_geom.hpp so the
// device assembler and this driver share the type without a circular include.
 
template <int Dim, unsigned Bits, class FluidFn>
FaceGeom buildFaceGeom(const AmrPoisson<Dim, Bits>& ap, FluidFn&& isFluid) {
  const Index n = ap.octree().numLeaves();
  std::vector<Index> start(static_cast<std::size_t>(n) + 1, 0);
  for (Index i = 0; i < n; ++i) {
    Index cnt = 0;
    ap.forEachFaceFull(i, [&](Index, int, int, double, double, double) { ++cnt; });
    start[static_cast<std::size_t>(i) + 1] = start[static_cast<std::size_t>(i)] + cnt;
  }
  const Index nf = start[static_cast<std::size_t>(n)];
  std::vector<Index> nbr(static_cast<std::size_t>(nf));
  std::vector<int> axis(static_cast<std::size_t>(nf)), dir(static_cast<std::size_t>(nf));
  std::vector<double> aArea(static_cast<std::size_t>(nf)), rArea(static_cast<std::size_t>(nf)),
      dist(static_cast<std::size_t>(nf)), alpha(static_cast<std::size_t>(nf));
  std::vector<Index> upupI(static_cast<std::size_t>(nf)), upupJ(static_cast<std::size_t>(nf));
  std::vector<double> invVol(static_cast<std::size_t>(n));
  std::vector<char> fluid(static_cast<std::size_t>(n));
  for (Index i = 0; i < n; ++i) {
    invVol[static_cast<std::size_t>(i)] = 1.0 / ap.cellVolume(i);
    fluid[static_cast<std::size_t>(i)] = isFluid(i) ? 1 : 0;
    Index k = start[static_cast<std::size_t>(i)];
    ap.forEachFaceFull(i, [&](Index j, int ax, int dr, double area, double d, double al) {
      nbr[static_cast<std::size_t>(k)] = j;
      axis[static_cast<std::size_t>(k)] = ax;
      dir[static_cast<std::size_t>(k)] = dr;
      aArea[static_cast<std::size_t>(k)] = al * area;
      rArea[static_cast<std::size_t>(k)] = area;
      dist[static_cast<std::size_t>(k)] = d;
      alpha[static_cast<std::size_t>(k)] = al;
      // SOU upstream-of-upwind probes (point neighbour one cell further upstream): if i is the
      // upwind cell the upstream is across i's −dir face; if j is upwind, across j's +dir face.
      upupI[static_cast<std::size_t>(k)] = ap.periodicNeighbor(i, ax, -dr);
      upupJ[static_cast<std::size_t>(k)] = ap.periodicNeighbor(j, ax, dr);
      ++k;
    });
  }
  FaceGeom g;
  g.n = n;
  g.start = toDevice(start, "fg_start");
  g.nbr = toDevice(nbr, "fg_nbr");
  g.axis = toDevice(axis, "fg_axis");
  g.dir = toDevice(dir, "fg_dir");
  g.alphaArea = toDevice(aArea, "fg_aarea");
  g.rawArea = toDevice(rArea, "fg_rarea");
  g.dist = toDevice(dist, "fg_dist");
  g.alpha = toDevice(alpha, "fg_alpha");
  g.upupI = toDevice(upupI, "fg_upupi");
  g.upupJ = toDevice(upupJ, "fg_upupj");
  g.invVol = toDevice(invVol, "fg_invvol");
  g.fluid = toDevice(fluid, "fg_fluid");
  return g;
}
 
inline void divergence(const FaceGeom& g, View<const double> u0, View<const double> u1,
                       View<const double> u2, View<double> div) {
  auto st = g.start;
  auto nb = g.nbr;
  auto ax = g.axis;
  auto dr = g.dir;
  auto aA = g.alphaArea;
  auto iv = g.invVol;
  auto fl = g.fluid;
  Kokkos::parallel_for(
      "amr::flow_div", g.n, KOKKOS_LAMBDA(const Index i) {
        if (!fl(i)) {
          div(i) = 0.0;
          return;
        }
        double d = 0.0;
        for (Index k = st(i); k < st(i + 1); ++k) {
          const int a = ax(k);
          const double ui = (a == 0) ? u0(i) : (a == 1) ? u1(i) : u2(i);
          const Index j = nb(k);
          const double uj = (a == 0) ? u0(j) : (a == 1) ? u1(j) : u2(j);
          d += aA(k) * dr(k) * 0.5 * (ui + uj);
        }
        div(i) = d * iv(i);
      });
}
 
inline void buildFaceField(const FaceGeom& g, View<const double> u0, View<const double> u1,
                           View<const double> u2, View<const double> phi, View<double> uf) {
  auto st = g.start;
  auto nb = g.nbr;
  auto ax = g.axis;
  auto dr = g.dir;
  auto di = g.dist;
  Kokkos::parallel_for(
      "amr::flow_buildface", g.n, KOKKOS_LAMBDA(const Index i) {
        for (Index k = st(i); k < st(i + 1); ++k) {
          const int a = ax(k);
          const Index j = nb(k);
          const double ui = (a == 0) ? u0(i) : (a == 1) ? u1(i) : u2(i);
          const double uj = (a == 0) ? u0(j) : (a == 1) ? u1(j) : u2(j);
          const double gphi = (dr(k) > 0) ? (phi(j) - phi(i)) / di(k) : (phi(i) - phi(j)) / di(k);
          uf(k) = 0.5 * (ui + uj) - gphi;
        }
      });
}
 
inline double divFaceNorm(const FaceGeom& g, View<const double> uf) {
  auto st = g.start;
  auto dr = g.dir;
  auto aA = g.alphaArea;
  auto iv = g.invVol;
  auto fl = g.fluid;
  double s = 0.0;
  Kokkos::parallel_reduce(
      "amr::flow_divface", g.n,
      KOKKOS_LAMBDA(const Index i, double& acc) {
        if (!fl(i))
          return;
        double d = 0.0;
        for (Index k = st(i); k < st(i + 1); ++k)
          d += aA(k) * dr(k) * uf(k);
        d *= iv(i);
        acc += d * d;
      },
      s);
  return std::sqrt(s);
}
 
inline void grad3(const FaceGeom& g, View<const double> f, View<double> gx, View<double> gy,
                  View<double> gz) {
  auto st = g.start;
  auto nb = g.nbr;
  auto ax = g.axis;
  auto dr = g.dir;
  auto di = g.dist;
  auto al = g.alpha;
  auto fl = g.fluid;
  Kokkos::parallel_for(
      "amr::flow_grad3", g.n, KOKKOS_LAMBDA(const Index i) {
        if (!fl(i)) {
          gx(i) = gy(i) = gz(i) = 0.0;
          return;
        }
        const double fi = f(i);
        double gp[3] = {0, 0, 0}, gm[3] = {0, 0, 0};
        int np[3] = {0, 0, 0}, nm[3] = {0, 0, 0};
        for (Index k = st(i); k < st(i + 1); ++k) {
          if (al(k) <= 1e-12)
            continue;
          const int a = ax(k);
          const double gg = (dr(k) > 0) ? (f(nb(k)) - fi) / di(k) : (fi - f(nb(k))) / di(k);
          if (dr(k) > 0) {
            gp[a] += gg;
            ++np[a];
          } else {
            gm[a] += gg;
            ++nm[a];
          }
        }
        double out[3];
        for (int a = 0; a < 3; ++a) {
          const double a1 = np[a] ? gp[a] / np[a] : 0.0;
          const double a2 = nm[a] ? gm[a] / nm[a] : 0.0;
          out[a] = 0.5 * (a1 + a2);
        }
        gx(i) = out[0];
        gy(i) = out[1];
        gz(i) = out[2];
      });
}
 
inline void momRhs(View<const double> uc, View<const double> gradP, View<const double> adv,
                   View<const double> rscale, View<const char> fluid, double idiag, double fc,
                   View<double> b, Index n) {
  Kokkos::parallel_for(
      "amr::flow_momrhs", n, KOKKOS_LAMBDA(const Index i) {
        b(i) = fluid(i) ? (idiag * uc(i) + fc - gradP(i) - adv(i)) * rscale(i) : 0.0;
      });
}
 
inline void buildFou(const FaceGeom& g, View<const double> u0, View<const double> u1,
                     View<const double> u2, double rho, View<const double> rscale,
                     View<double> advDiag, View<double> advCoef, View<const double> uf,
                     bool useFace) {
  auto st = g.start;
  auto nb = g.nbr;
  auto ax = g.axis;
  auto dr = g.dir;
  auto ra = g.rawArea;
  auto iv = g.invVol;
  auto fl = g.fluid;
  Kokkos::parallel_for(
      "amr::flow_buildfou", g.n, KOKKOS_LAMBDA(const Index i) {
        if (!fl(i)) {
          advDiag(i) = 0.0;
          for (Index k = st(i); k < st(i + 1); ++k)
            advCoef(k) = 0.0;
          return;
        }
        const double rs = rscale(i);
        double dsum = 0.0;
        for (Index k = st(i); k < st(i + 1); ++k) {
          const Index j = nb(k);
          if (!fl(j)) {
            advCoef(k) = 0.0;
            continue;
          }
          const int a = ax(k);
          const double ui = (a == 0) ? u0(i) : (a == 1) ? u1(i) : u2(i);
          const double uj = (a == 0) ? u0(j) : (a == 1) ? u1(j) : u2(j);
          const double velOut = useFace ? dr(k) * uf(k) : dr(k) * 0.5 * (ui + uj);
          const double w = rs * rho * ra(k) * velOut * iv(i);
          if (velOut < 0.0) {
            advCoef(k) = w;  // inflow → off-diagonal toward upstream neighbour j
          } else {
            advCoef(k) = 0.0;
            dsum += w;  // outflow → diagonal
          }
        }
        advDiag(i) = dsum;
      });
}
 
inline void deferredSou(const FaceGeom& g, View<const double> u0, View<const double> u1,
                        View<const double> u2, int comp, double rho, int advScheme,
                        View<double> defc, View<const double> uf, bool useFace) {
  auto st = g.start;
  auto nb = g.nbr;
  auto ax = g.axis;
  auto dr = g.dir;
  auto ra = g.rawArea;
  auto iv = g.invVol;
  auto fl = g.fluid;
  auto uiP = g.upupI;
  auto ujP = g.upupJ;
  Kokkos::parallel_for(
      "amr::flow_defsou", g.n, KOKKOS_LAMBDA(const Index i) {
        if (!fl(i)) {
          defc(i) = 0.0;
          return;
        }
        auto fld = [&](Index c) { return (comp == 0) ? u0(c) : (comp == 1) ? u1(c) : u2(c); };
        double sou = 0.0, fou = 0.0;
        for (Index k = st(i); k < st(i + 1); ++k) {
          const Index j = nb(k);
          if (!fl(j))
            continue;
          const int a = ax(k);
          const double uai = (a == 0) ? u0(i) : (a == 1) ? u1(i) : u2(i);
          const double uaj = (a == 0) ? u0(j) : (a == 1) ? u1(j) : u2(j);
          const double velOut = useFace ? dr(k) * uf(k) : dr(k) * 0.5 * (uai + uaj);
          const Index up = (velOut > 0.0) ? i : j;
          const Index down = (velOut > 0.0) ? j : i;
          const Index upup = (velOut > 0.0) ? uiP(k) : ujP(k);
          const double phiUp = fld(up);
          const double phiUpUp = (upup >= 0 && fl(upup)) ? fld(upup) : phiUp;
          const double phiDown = fld(down);
          const double phiFace = hoFaceValue(phiUpUp, phiUp, phiDown, advScheme);  // shared recon
          sou += ra(k) * velOut * phiFace;
          fou += ra(k) * velOut * fld(up);  // FOU flux = velOut · upwind value
        }
        defc(i) = rho * iv(i) * (sou - fou);  // ρ·(SOU − FOU), unscaled
      });
}
 
inline void advectExplicit(const FaceGeom& g, View<const double> u0, View<const double> u1,
                           View<const double> u2, int comp, double rho, int advScheme,
                           View<double> defc, View<const double> uf, bool useFace) {
  auto st = g.start;
  auto nb = g.nbr;
  auto ax = g.axis;
  auto dr = g.dir;
  auto ra = g.rawArea;
  auto iv = g.invVol;
  auto fl = g.fluid;
  auto uiP = g.upupI;
  auto ujP = g.upupJ;
  Kokkos::parallel_for(
      "amr::flow_advexpl", g.n, KOKKOS_LAMBDA(const Index i) {
        if (!fl(i)) {
          defc(i) = 0.0;
          return;
        }
        auto fld = [&](Index c) { return (comp == 0) ? u0(c) : (comp == 1) ? u1(c) : u2(c); };
        double sou = 0.0;
        for (Index k = st(i); k < st(i + 1); ++k) {
          const Index j = nb(k);
          if (!fl(j))
            continue;
          const int a = ax(k);
          const double uai = (a == 0) ? u0(i) : (a == 1) ? u1(i) : u2(i);
          const double uaj = (a == 0) ? u0(j) : (a == 1) ? u1(j) : u2(j);
          const double velOut = useFace ? dr(k) * uf(k) : dr(k) * 0.5 * (uai + uaj);
          const Index up = (velOut > 0.0) ? i : j;
          const Index down = (velOut > 0.0) ? j : i;
          const Index upup = (velOut > 0.0) ? uiP(k) : ujP(k);
          const double phiUp = fld(up);
          const double phiUpUp = (upup >= 0 && fl(upup)) ? fld(upup) : phiUp;
          const double phiDown = fld(down);
          const double phiFace = hoFaceValue(phiUpUp, phiUp, phiDown, advScheme);  // shared recon
          sou += ra(k) * velOut * phiFace;
        }
        defc(i) = rho * sou * iv(i);  // ρ·SOU (fully explicit)
      });
}
 
inline void correct(View<double> uc, View<const double> gphi, View<const char> fluid, Index n) {
  Kokkos::parallel_for(
      "amr::flow_correct", n, KOKKOS_LAMBDA(const Index i) {
        if (fluid(i))
          uc(i) -= gphi(i);
      });
}
 
inline void presUpdate(View<double> p, View<const double> phi, View<const double> div,
                       View<const char> fluid, double rho_dt, double mu, Index n) {
  Kokkos::parallel_for(
      "amr::flow_presupd", n, KOKKOS_LAMBDA(const Index i) {
        if (fluid(i))
          p(i) += rho_dt * phi(i) - mu * div(i);
      });
}
 
// ===========================================================================
// AmrFlow — collocated Stokes projection step, fully on device.
// ===========================================================================
template <unsigned Bits = 21u>
class AmrFlow {
 public:
  using Octree = BlockOctree<3, Bits>;
 
  void init(const Octree& t, Real h0, Vec<3> origin = Vec<3>{}) {
    t_ = &t;
    h0_ = h0;
    origin_ = origin;
  }
  void setDensity(double rho) { rho_ = rho; }
  void setViscosity(double mu) { mu_ = mu; }
  void setDt(double dt) { dt_ = dt; }
  void setBodyForce(double fx, double fy, double fz) { f_ = {fx, fy, fz}; }
  void setPressurePCG(bool on) { presPCG_ = on; }
  void setAdvection(bool on) { advect_ = on; }
  void setImplicitAdvection(bool on) { implicitFou_ = on; }
  void setAdvectionScheme(int s) { advScheme_ = s; }
  void setMomentumTol(double tol) { momTol_ = tol; }
  void setMomentumMG(bool on) { momMGon_ = on; }
 
  void setVelocityMGStaircase(bool on) { useStaircaseMG_ = on; }
  void setVelocityMGMinCoarse(Index m) { mgMinCoarse_ = m; }
  void setMomentumGS(bool on) { momGS_ = on; }
 
  void setMomentumMGSolver(bool on) { momMGSolver_ = on; }
 
  void setOuterIterations(int n, double tol = 1e-6) {
    outerIters_ = (n < 1) ? 1 : n;
    outerTol_ = tol;
  }
 
  template <class SdfFn>
  void setSolid(SdfFn&& sdfFn) {
    const Index n = t_->numLeaves();
    // Host operator build (geometry, openness, cut stencils) — same as oracle::AmrFlow::setSolid.
    mom_.init(*t_, h0_, origin_);
    mom_.build(sdfFn, /*idiag=*/rho_ / dt_, /*beta=*/mu_ / (h0_ * h0_));
    pres_.init(*t_, h0_);
    pres_.setOrigin(origin_);
    auto openFn = [&](const Vec<3>& fc, int axis) { return faceFrac(sdfFn, fc, axis); };
    pres_.buildOpenness(openFn);
    presMG_.build(*t_, h0_, openFn, /*periodic=*/true);
    presMG_.setRemoveMean(true);  // singular periodic pressure: per-level nullspace projection
 
    // Device assembly (D6): the static cut-cell momentum operator CSR and the collocated face
    // geometry are assembled ON THE DEVICE (assembleMomentum / assembleFaceGeom), and the
    // pressure MG operators are device-assembled per level (Multigrid D5) — so no host CSR walk and
    // no operator round-trip. Each is bit-for-bit identical to the host assembler on OpenMP (locked
    // in test_amr_momentum / test_amr_facegeom), so the flow result is unchanged. A shared device
    // octree view backs both assemblers.
    BlockOctreeView<3, Bits> ov;
    ov.upload(*t_);
    momOp_ = assembleMomentum<Bits>(mom_, ov);  // static Stokes operator (hasAdv stays false)
    // Velocity multigrid (momentum preconditioner): the Galerkin hierarchy A_c = R·A·P built
    // directly from the exact assembled momentum CSR. Consistent with the fine cut-cell
    // operator by construction (inherits the ξ-overlay + D_rescale row scaling; a coarse cell
    // of all-solid children stays an identity row). It only changes the preconditioner (the
    // BiCGStab matvec is the exact operator) ⇒ same converged solution, but the iteration
    // count stays ~flat with N instead of growing like the Jacobi-preconditioned BiCGStab.
    // Build the chosen momentum-MG: Galerkin (MomentumMG) by default, or the rediscretized
    // staircase (VelocityMG) — both from the static Stokes operator, once.
    if (momMGon_) {
      if (useStaircaseMG_) {
        std::vector<double> kap(static_cast<std::size_t>(n));
        std::vector<char> fl(static_cast<std::size_t>(n)), cu(static_cast<std::size_t>(n));
        for (Index i = 0; i < n; ++i) {
          kap[static_cast<std::size_t>(i)] = mom_.kappa(i);
          fl[static_cast<std::size_t>(i)] = mom_.isFluid(i) ? 1 : 0;
          cu[static_cast<std::size_t>(i)] = mom_.isCut(i) ? 1 : 0;
        }
        velMG_.build(*t_, h0_, rho_ / dt_, mu_, momOp_, kap, fl, cu, mgMinCoarse_);
        velMG_.setGaussSeidel(momGS_);
      } else {
        // The Galerkin RAP hierarchy is a host triple-product over the fine CSR, so it needs the
        // operator on the host; assemble it there for the MG build only (bit-identical to momOp_).
        auto A = mom_.assembleOperator();
        momMG_.build(*t_, A.diag, A.start, A.nbr, A.coef);
        momMG_.setGaussSeidel(momGS_);
      }
    }
    std::vector<char> fluidVec(static_cast<std::size_t>(n));
    for (Index i = 0; i < n; ++i)
      fluidVec[static_cast<std::size_t>(i)] = mom_.isFluid(i) ? 1 : 0;
    geom_ = assembleFaceGeom<Bits>(pres_, fluidVec, ov);
    std::vector<double> rs(static_cast<std::size_t>(n));
    for (Index i = 0; i < n; ++i)
      rs[static_cast<std::size_t>(i)] = mom_.rhsScale(i);
    rscale_ = toDevice(rs, "df_rscale");
    fluid_ = geom_.fluid;
 
    // Device state.
    for (int c = 0; c < 3; ++c) {
      u_[c] = View<double>("df_u", static_cast<std::size_t>(n));
      gx_[c] = View<double>("df_g", static_cast<std::size_t>(n));
      Kokkos::deep_copy(u_[c], 0.0);
    }
    p_ = View<double>("df_p", static_cast<std::size_t>(n));
    phi_ = View<double>("df_phi", static_cast<std::size_t>(n));
    div_ = View<double>("df_div", static_cast<std::size_t>(n));
    uf_ = View<double>("df_uf",
                       geom_.nbr.extent(0));  // ABC divergence-free face field (per CSR face)
    faceFieldBuilt_ = false;
    bmom_ = View<double>("df_bmom", static_cast<std::size_t>(n));
    Kokkos::deep_copy(p_, 0.0);
    // Implicit-FOU advection state. The momentum operator + its velocity-MG are rebuilt each
    // step from the *full* operator (viscous + FOU) so the MG is advection-aware (the viscous-
    // only MG diverges on the advection operator at cut cells); the FOU is baked into the CSR
    // (momOp_.hasAdv stays false). defc holds the device-computed explicit ρ(SOU−FOU)
    // deferred correction. uadvHost_ caches u^n on the host for the per-step operator rebuild.
    for (int c = 0; c < 3; ++c) {
      defc_[c] = View<double>("df_defc", static_cast<std::size_t>(n));
      Kokkos::deep_copy(defc_[c], 0.0);
      // uⁿ snapshot for the backward-Euler mass term (frozen across Picard outer iters) + the
      // previous outer iterate for the outer-loop convergence test.
      u0_[c] = View<double>("df_u0", static_cast<std::size_t>(n));
      uprev_[c] = View<double>("df_uprev", static_cast<std::size_t>(n));
    }
    // Device-resident implicit-FOU advection: the FOU operator (advDiag + per-face advCoef over the
    // face-geometry CSR) is rebuilt on device each step from uⁿ and added to the static Stokes
    // operator in the matvec (no host round-trip). The static operator + Galerkin velocity-MG are
    // built once (above); the advection is a perturbation the static MG still preconditions.
    advDiag_ = View<double>("df_advdiag", static_cast<std::size_t>(n));
    advCoef_ = View<double>("df_advcoef", geom_.nbr.extent(0));
    momOp_.advStart = geom_.start;
    momOp_.advNbr = geom_.nbr;
    momOp_.advDiag = advDiag_;
    momOp_.advCoef = advCoef_;
    momOp_.hasAdv = false;  // set per-step when advection is on
    momSolver_.setJacobi(2, 0.7);
    // Generic MG preconditioner: dispatch the chosen hierarchy's V-cycle (z = M⁻¹ r), decoupling
    // the solver from the MG type so Galerkin and staircase are interchangeable.
    if (momMGon_) {
      if (useStaircaseMG_)
        momSolver_.setPreconditioner(
            [this](View<const double> r, View<double> z) { runMgVcycle(velMG_, r, z); });
      else
        momSolver_.setPreconditioner(
            [this](View<const double> r, View<double> z) { runMgVcycle(momMG_, r, z); });
    }
    pcg_.setVcycle(2, 2, 60, 0.8);
    pcg_.setSingular(true);
    n_ = n;
  }
 
  void step(int momIters = 100, int presIters = 60) {
    const Index n = n_;
    const double idiag = rho_ / dt_;
    lastMomIters_ = 0;
    lastOuterIters_ = 1;
    // Freeze the time-level uⁿ for the backward-Euler mass term + warm start; it stays anchored
    // across the Picard outer iterations (only the advecting velocity re-lags). For outerIters_==1
    // this is just a copy of uⁿ ⇒ bit-identical to the single lagged step.
    for (int c = 0; c < 3; ++c)
      Kokkos::deep_copy(u0_[c], View<const double>(u_[c]));
    // −∇p^n is constant across the outer iterations (pressure is projected once, after the loop,
    // like flow's single per-step projection) ⇒ hoist it out.
    grad3(geom_, View<const double>(p_), gx_[0], gx_[1], gx_[2]);
    // Picard outer loop over the lagged advection only (the momentum nonlinearity); for
    // outerIters_==1 this is the single lagged predictor, then one projection — bit-identical to
    // before.
    for (int outer = 0; outer < outerIters_; ++outer) {
      // --- advection lagged to the *current* predictor iterate (uⁿ on the first pass):
      // implicit-FOU operator + explicit ρ(SOU−FOU) deferred correction. The matvec stays linear
      // during each solve (the advecting velocity is frozen in advDiag_/advCoef_). With advection
      // OFF the operator and RHS are identical every pass, so a second pass reproduces the first ⇒
      // instant early-stop. ---
      if (advect_) {
        momOp_.hasAdv = implicitFou_;
        // Advect with the divergence-free face field uf (from the previous projection); fall back
        // to ½(u_i+u_j) before the first projection has built it. The implicit FOU and the explicit
        // SOU/FOU deferred correction use the SAME velocity ⇒ the FOU cancels at steady state
        // (host-parity).
        const View<const double> ufv(uf_);
        if (implicitFou_)
          buildFou(geom_, View<const double>(u_[0]), View<const double>(u_[1]),
                   View<const double>(u_[2]), rho_, View<const double>(rscale_), advDiag_, advCoef_,
                   ufv, faceFieldBuilt_);
        for (int c = 0; c < 3; ++c) {
          if (implicitFou_)
            deferredSou(geom_, View<const double>(u_[0]), View<const double>(u_[1]),
                        View<const double>(u_[2]), c, rho_, advScheme_, defc_[c], ufv,
                        faceFieldBuilt_);
          else
            advectExplicit(geom_, View<const double>(u_[0]), View<const double>(u_[1]),
                           View<const double>(u_[2]), c, rho_, advScheme_, defc_[c], ufv,
                           faceFieldBuilt_);
        }
        // The staircase MG's fine level mirrors the sharp operator; refresh it so it picks up the
        // current advection state (hasAdv). (The Galerkin MG is the static viscous operator.)
        if (momMGon_ && useStaircaseMG_)
          velMG_.setFineOp(momOp_);
      }
      // --- predictor: incremental BE viscous (+ implicit-FOU) solve per component, RHS carries
      // −∇p^n and −ρ(SOU−FOU); the mass term is anchored at uⁿ (u0_), the solve warm-starts at the
      // current iterate. ---
      for (int c = 0; c < 3; ++c) {
        momRhs(View<const double>(u0_[c]), View<const double>(gx_[c]), View<const double>(defc_[c]),
               View<const double>(rscale_), View<const char>(fluid_), idiag, f_[c], bmom_, n);
        // P4 (opt-in): the velocity-MG used as the *solver* — MG-preconditioned defect correction,
        // no Krylov (the flow RB-GS/velocity-MG mirror; cannot break down on the non-symmetric
        // operator) — vs the default MG-preconditioned BiCGStab. Both reach the same solution (the
        // matvec is the exact operator); the choice only trades robustness for convergence rate.
        lastMomIters_ +=
            (momMGSolver_ ? momSolver_.solveDefectCorrection(
                                momOp_, u_[c], View<const double>(bmom_), momIters, momTol_)
                          : momSolver_.solveBiCGStab(momOp_, u_[c], View<const double>(bmom_),
                                                     momIters, momTol_))
                .iters;
      }
      // Outer-loop convergence on the predictor velocity (skipped for the default outerIters_==1,
      // so that path is untouched). With advection off the second iterate equals the first ⇒ stops
      // at 2.
      if (outerIters_ > 1) {
        lastOuterIters_ = outer + 1;
        if (outer > 0) {
          double dmax = 0.0;
          for (int c = 0; c < 3; ++c)
            dmax = std::max(
                dmax, maxAbsDiff(View<const double>(u_[c]), View<const double>(uprev_[c]), n));
          if (dmax < outerTol_)
            break;
        }
        for (int c = 0; c < 3; ++c)
          Kokkos::deep_copy(uprev_[c], View<const double>(u_[c]));
      }
    }
    project(presIters);  // single pressure projection per step (flow structure)
  }
 
  void project(int presIters = 60) {
    const Index n = n_;
    divergence(geom_, View<const double>(u_[0]), View<const double>(u_[1]),
               View<const double>(u_[2]), div_);
    Kokkos::deep_copy(phi_, 0.0);
    // Two selectable pressure drivers, like flow's CutcellMG (MG-PCG single-rank default,
    // standalone V-cycle the robust multi-rank default): MG-PCG for the near-steady, divergence-
    // compatible Stokes case (faster), and the stationary V-cycle for the large transient
    // divergence of advection, which excites near-nullspace modes of the cut-cell openness Poisson
    // (near-isolated fluid cells) that CG amplifies — the V-cycle is bounded. (Per-level mean
    // removal is on for both; making MG-PCG robust enough to also cover advection needs flow's
    // near-isolated-cell pinning/classification — a follow-up.)
    if (presPCG_ && !advect_) {
      lastPresIters_ = pcg_.solve(presMG_, phi_, View<const double>(div_), presIters, 1e-10).iters;
    } else {
      Kokkos::deep_copy(presMG_.b(0), div_);
      Kokkos::deep_copy(presMG_.x(0), 0.0);
      for (int it = 0; it < presIters; ++it)
        presMG_.vcycle(2, 2, 60, 0.8);
      Kokkos::deep_copy(phi_, presMG_.x(0));
      lastPresIters_ = presIters;
    }
    // Build the divergence-free FACE field from u* (still in u_) + φ, before the cell correction.
    buildFaceField(geom_, View<const double>(u_[0]), View<const double>(u_[1]),
                   View<const double>(u_[2]), View<const double>(phi_), uf_);
    faceFieldBuilt_ = true;
    grad3(geom_, View<const double>(phi_), gx_[0], gx_[1], gx_[2]);
    for (int c = 0; c < 3; ++c)
      correct(u_[c], View<const double>(gx_[c]), View<const char>(fluid_), n);
    presUpdate(p_, View<const double>(phi_), View<const double>(div_), View<const char>(fluid_),
               rho_ / dt_, mu_, n);
  }
 
  std::vector<double> debugSou(int comp) {
    View<double> s("dbg_sou", static_cast<std::size_t>(n_));
    advectExplicit(geom_, View<const double>(u_[0]), View<const double>(u_[1]),
                   View<const double>(u_[2]), comp, 1.0, advScheme_, s, View<const double>(uf_),
                   false);
    std::vector<double> h(static_cast<std::size_t>(n_));
    auto m = Kokkos::create_mirror_view(s);
    Kokkos::deep_copy(m, s);
    for (Index i = 0; i < n_; ++i)
      h[static_cast<std::size_t>(i)] = m(i);
    return h;
  }
  template <class MG>
  void runMgVcycle(MG& mg, View<const double> r, View<double> z) {
    Kokkos::deep_copy(mg.b(0), r);
    Kokkos::deep_copy(mg.x(0), 0.0);
    mg.vcycle(mgVcPre_, mgVcPre_, mgVcBottom_, 0.7);
    Kokkos::deep_copy(z, mg.x(0));
  }
  static double maxAbsDiff(View<const double> a, View<const double> b, Index n) {
    double m = 0.0;
    Kokkos::parallel_reduce(
        "amr::flow_maxdiff", n,
        KOKKOS_LAMBDA(const Index i, double& lm) {
          double d = a(i) - b(i);
          if (d < 0.0)
            d = -d;
          if (d > lm)
            lm = d;
        },
        Kokkos::Max<double>(m));
    return m;
  }
  void copyToHost(const View<double>& d, std::vector<double>& h) const {
    auto m = Kokkos::create_mirror_view(d);
    Kokkos::deep_copy(m, d);
    for (Index i = 0; i < n_; ++i)
      h[static_cast<std::size_t>(i)] = m(i);
  }
  void setVelocity(int c, const std::vector<double>& h) {
    auto m = Kokkos::create_mirror_view(u_[c]);
    for (Index i = 0; i < n_; ++i)
      m(i) = h[static_cast<std::size_t>(i)];
    Kokkos::deep_copy(u_[c], m);
  }
  std::vector<double> velocity(int c) const { return peclet::core::toVector(u_[c]); }
  std::vector<double> pressure() const { return peclet::core::toVector(p_); }
 
  std::vector<double> velocities() const {
    View<double> packed(Kokkos::view_alloc("amr::vel_packed", Kokkos::WithoutInitializing),
                        static_cast<std::size_t>(n_) * 3);
    for (int c = 0; c < 3; ++c) {
      auto uc = u_[c];
      auto p = packed;
      const int cc = c;
      Kokkos::parallel_for(
          "amr::pack_vel", n_, KOKKOS_LAMBDA(const Index i) { p(i * 3 + cc) = uc(i); });
    }
    return peclet::core::toVector(packed);
  }
  double divNormL2() {
    divergence(geom_, View<const double>(u_[0]), View<const double>(u_[1]),
               View<const double>(u_[2]), div_);
    return std::sqrt(dotPlain(View<const double>(div_), View<const double>(div_), n_));
  }
  double divNormFace() { return divFaceNorm(geom_, View<const double>(uf_)); }
  std::vector<double> faceField() const { return peclet::core::toVector(uf_); }
  Index numLeaves() const { return n_; }
  bool isFluid(Index i) const { return mom_.isFluid(i); }
  int lastMomIters() const { return lastMomIters_; }
  int lastPresIters() const { return lastPresIters_; }
  int lastOuterIters() const { return lastOuterIters_; }
 
 private:
  // flow ccFractionCore aperture (verbatim from oracle::AmrFlow::faceFrac).
  template <class SdfFn>
  double faceFrac(SdfFn&& sdfFn, const Vec<3>& fc, int axis) const {
    double sd = sdfFn(fc);
    if (sd <= 0.0)
      return 0.0;
    Vec<3> g{};
    for (int d = 0; d < 3; ++d) {
      Vec<3> pp = fc, pm = fc;
      pp[d] += h0_;
      pm[d] -= h0_;
      g[d] = (sdfFn(pp) - sdfFn(pm)) / (2.0 * h0_);
    }
    double gmag = std::sqrt(g[0] * g[0] + g[1] * g[1] + g[2] * g[2]);
    if (gmag < 1e-6)
      gmag = 1e-6;
    int t1 = (axis + 1) % 3, t2 = (axis + 2) % 3;
    double denom = (std::fabs(g[t1]) + std::fabs(g[t2])) / gmag * h0_;
    if (denom < 1e-9)
      denom = 1e-9;
    double frac = 0.5 + sd / denom;
    return frac < 0.0 ? 0.0 : (frac > 1.0 ? 1.0 : frac);
  }
 
  const Octree* t_ = nullptr;
  Real h0_ = 1.0;
  Vec<3> origin_{};
  double rho_ = 1.0, mu_ = 1.0, dt_ = 1e6;
  Vec<3> f_{};
  bool presPCG_ = true;
  bool momMGon_ = true;  // velocity-MG momentum preconditioner (scalable; see setMomentumMG)
  bool useStaircaseMG_ = false;  // false = Galerkin (MomentumMG), true = staircase (VelocityMG)
  int mgVcPre_ = 2, mgVcBottom_ = 30;  // momentum-MG V-cycle pre/post sweeps + bottom sweeps
  Index mgMinCoarse_ = 256;            // staircase velocity-MG pore-scale cap (coarsest cell count)
  bool momGS_ = false;  // opt-in: multicolour Gauss–Seidel smoother in the momentum MG
  bool momMGSolver_ =
      false;            // opt-in (P4): velocity-MG as the solver (defect correction), not BiCGStab
  int outerIters_ = 1;  // Picard outer iterations over the lagged advection (default 1)
  double outerTol_ = 1e-6;   // outer-loop early-stop tolerance on max|Δu|
  double momTol_ = 1e-8;     // per-step momentum BiCGStab relative tolerance (Phase-0 knob)
  bool advect_ = false;      // momentum advection ∇·(u u) (off ⇒ Stokes)
  bool implicitFou_ = true;  // implicit-FOU deferred-correction (stable) vs fully-explicit
  int advScheme_ = 0;        // high-order flux: 0 = SOU (default), 1 = Koren TVD
  Index n_ = 0;
  int lastMomIters_ = 0, lastPresIters_ = 0, lastOuterIters_ = 1;
 
  AmrCutCell<Bits> mom_;
  AmrPoisson<3, Bits> pres_;
  Multigrid<3, Bits> presMG_;
  MomentumMG<Bits> momMG_;  // Galerkin velocity multigrid (momentum preconditioner)
  VelocityMG<Bits> velMG_;  // rediscretized staircase velocity multigrid (alternative)
  MomentumOp momOp_;
  MomentumSolver<Bits> momSolver_;
  PCG<3, Bits> pcg_;
  std::array<View<double>, 3> defc_;  // explicit ρ(SOU−FOU) deferred correction per component
  View<double> advDiag_, advCoef_;    // device-resident implicit-FOU operator (rebuilt each step)
  FaceGeom geom_;
  View<double> rscale_;
  View<char> fluid_;
  std::array<View<double>, 3> u_, gx_;
  std::array<View<double>, 3> u0_,
      uprev_;  // frozen uⁿ (BE mass term) + previous Picard outer iterate
  View<double> p_, phi_, div_, bmom_;
  View<double> uf_;  // ABC/Basilisk divergence-free face field (one per CSR (sub)face)
  bool faceFieldBuilt_ =
      false;  // uf_ populated by a projection (else advection falls back to ½(u_i+u_j))
};
 
}  // namespace peclet::core::amr
 
#endif  // PECLET_CORE_HAVE_MORTON
#endif  // PECLET_CORE_AMR_FLOW_HPP