flow_bindings.cpp

Go to the documentation of this file.

 
#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
#include <nanobind/stl/pair.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/vector.h>
 
#include <array>
#include <cstdint>
#include <Kokkos_Core.hpp>
#include <string>
#include <vector>
 
#ifdef PECLET_FLOW_MPI
#include <mpi.h>
 
#include "peclet/core/common/types.hpp"
#include "peclet/core/decomp/block_decomposer.hpp"
#endif
 
#include "flow_ibm.hpp"
#include "peclet/core/python/ndarray_interop.hpp"
 
#ifdef PECLET_FLOW_MPI
// Ensure MPI_Init has been called (mirrors the dem init_mpi idiom); safe to call repeatedly.
static void ensure_mpi_init() {
  int inited = 0;
  MPI_Initialized(&inited);
  if (!inited) {
    int argc = 0;
    char** argv = nullptr;
    MPI_Init(&argc, &argv);
  }
}
#endif
 
namespace nb = nanobind;
 
// A solver field (flat x-fastest, ghost-stripped) -> Fortran-order (nx,ny,nz) float64 NumPy array
// for [x,y,z] indexing. The vector is moved into the array's backing store (no extra copy vs the
// old to_xyz).
template <class S>
static nb::ndarray<nb::numpy, double> field_out(S& s, std::vector<double>&& v) {
  const auto nx = static_cast<std::size_t>(s.nx());
  const auto ny = static_cast<std::size_t>(s.ny());
  const auto nz = static_cast<std::size_t>(s.nz());
  return peclet::core::python::vector_to_ndarray(
      std::move(v), {nx, ny, nz},
      {1, static_cast<std::int64_t>(nx), static_cast<std::int64_t>(nx * ny)});
}
 
// A Fortran-order (nx,ny,nz) float64 array -> flat x-fastest host vector (F-contiguous data() is
// already x-fastest). nanobind casts/copies the input to f_contig double if needed.
static std::vector<double> grid_in(nb::ndarray<double, nb::f_contig> a) {
  return peclet::core::python::ndarray_to_vector<double>(nb::ndarray<>(a));
}
 
// Zero-copy export of a registered field's padded device buffer as a Fortran-order 3-D array of the
// full block shape (ex,ey,ez) = (nx+2G, ny+2G, nz+2G), x-fastest strides {1,ex,ex*ey}. Includes the
// ghost band (the flat buffer is contiguous; a ghost-stripped view would not be). The capsule owns
// a copy of the managed CCField, so the allocation outlives the array — host → NumPy referencing
// the buffer, device → DLPack for CuPy/torch. Mirrors peclet::core::python::view_to_ndarray but
// with an explicit 3-D reshape of the flat 1-D field.
template <class S>
static auto field3d_out(S& s, peclet::flow::CCField f) {
  namespace pcp = peclet::core::python;
  using Mem = peclet::flow::CCMem;
  const auto bs = s.blockShape();
  std::array<std::size_t, 3> shape{static_cast<std::size_t>(bs[0]), static_cast<std::size_t>(bs[1]),
                                   static_cast<std::size_t>(bs[2])};
  std::array<std::int64_t, 3> strides{1, static_cast<std::int64_t>(bs[0]),
                                      static_cast<std::int64_t>(bs[0]) * bs[1]};
  auto* held = new peclet::flow::CCField(f);
  nb::capsule owner(held, [](void* p) noexcept { delete static_cast<peclet::flow::CCField*>(p); });
  double* data = f.data();
  if constexpr (pcp::is_host_space_v<Mem>) {
    return nb::ndarray<nb::numpy, double>(data, 3, shape.data(), owner, strides.data(),
                                          nb::dtype<double>(), nb::device::cpu::value, 0);
  } else {
    auto dev = pcp::dlpack_device<Mem>();
    return nb::ndarray<double>(data, 3, shape.data(), owner, strides.data(), nb::dtype<double>(),
                               dev.first, dev.second);
  }
}
 
// Register a solver class for the given GridLayout policy (Staggered -> "Solver", Colocated ->
// "SolverColocated"). The Python API is identical across grids; only the velocity-unknown placement
// and the advection control volume differ inside Solver<Grid>.
template <class Grid>
static void bind_solver(nb::module_& m, const char* name) {
  using S = peclet::flow::Solver<Grid>;
  nb::class_<S>(m, name)
      .def(nb::init<int, int, int>(), nb::arg("nx"), nb::arg("ny"), nb::arg("nz"),
           "Create a solver on an nx x ny x nz unit-spacing grid (x-fastest, I = x + y*nx + "
           "z*nx*ny). "
           "Set physical parameters (rho/mu/dt) and any domain BCs before the geometry / first "
           "step.")
      .def("set_rho", &S::setRho, nb::arg("rho"),
           "Set fluid density rho (physical units). Set before geometry/first step.")
      .def("set_mu", &S::setMu, nb::arg("mu"), "Set dynamic viscosity mu (physical units).")
      .def("set_dt", &S::setDt, nb::arg("dt"),
           "Set the time step dt; the momentum solve is scaled by 1/dt (well-conditioned at large "
           "dt).")
      .def("set_body_force", &S::setBodyForce, nb::arg("fx"), nb::arg("fy"), nb::arg("fz"),
           "Set the body force per unit volume (fx, fy, fz) — e.g. a mean pressure gradient.")
      .def("set_advection", &S::setAdvection, nb::arg("on"),
           "Enable/disable explicit high-order momentum advection (default scheme SOU). Off ⇒ "
           "Stokes.")
      .def("set_advection_scheme", &S::setAdvectionScheme, nb::arg("scheme"),
           "High-order advection scheme: 0 = second-order upwind (SOU, default), 1 = Koren TVD.")
      .def("set_incremental_pressure", &S::setIncrementalPressure, nb::arg("on"),
           "Toggle the rotational incremental-pressure projection.")
      .def("set_pressure_warmstart", &S::setPressureWarmstart, nb::arg("on"),
           "Seed each pressure solve from the previous step's phi (default off).")
      .def(
          "set_face_interp", &S::setFaceInterp, nb::arg("mode"),
          "DEPRECATED integer form of set_collocated_scheme (0 = plain, 9 = gauge-exact, the "
          "default). Modes 1/2/5/6/7/10 were RETIRED 2026-08-18 (ablations; 10 measured divergent) "
          "and now raise. Modes 3/4 survive as FV-constraint ablations (4 pairs with "
          "set_fv_relax). No effect on the staggered solver.")
      .def(
          "set_collocated_scheme", &S::setCollocatedScheme, nb::arg("name"),
          "Collocated cut-cell projection scheme (no effect on the staggered solver). DEFAULT "
          "since 2026-08-25: AUTO = 'ghost' where supported, falling back to 'gauge-exact' with "
          "a stderr notice on porous/variable-rho/domain-BC/Chebyshev configurations (ghost v1 "
          "limits); any explicit selection here disables AUTO.\n"
          "  'gauge-exact' (default 2026-08-18..25; the AUTO fallback) — the aperture constraint with the "
          "directional gauge-exact pressure gradient. Cheapest scheme measured (symmetric MG-PCG, "
          "no fragmentation guard). CAVEATS from the 2026-08 campaign "
          "(doc/collocated_invisible_subspace.md): possesses an attractor FAMILY of steady states "
          "(support-inconsistent gradient/constraint pair) selected by march protocol, and a "
          "rotational-update instability whose wall-blend stopgap "
          "(set_rotational_wall_weight) has a resolution-dependent margin — UNSTABLE at "
          "(R>=16 cells/radius, dt>=600); use fixed dt<=60 protocols at high resolution. "
          "Clean-protocol bias vs the staggered reference: -2.5% (R=8) -> ~+0.2% asymptote.\n"
          "  'ghost' — the fluid-only constraint scheme (binary-openness divergence + directional "
          "closures + the same gauge-exact gradient). Family-free, unconditionally stable "
          "(dt 60..1e20, no stabilizer), protocol-independent (C2), both-bed clean ladder "
          "-1.4% (R=8) -> +0.22% asymptote, Z&H anchor -0.018% at N=128. Costs: nonsymmetric "
          "BiCGStab pressure solve (~2.3-2.7x), ~1.6 KB/cell overlay (single-GPU size cap), "
          "fragmentation guard. MPI-capable (np=1,2,4 ctests); at-scale np>=16 hardening in "
          "progress. Equivalent to set_ghost_projection(True, 2, 2); the (1, 2) mixed mode "
          "stays quarantined.\n"
          "  'plain' — the legacy path: plain 1/2-1/2 face average + central-difference grad(P). "
          "FIRST order at curved cut cells (the cell gradient reads the decoupled p=0 of "
          "solid-centred neighbours, an O(1/h) gauge error), and on a dense bed it also fails to "
          "reach steady state within 800 steps at coarse resolution. Kept for reproducing "
          "published results.")
      .def(
          "set_rotational_pressure", &S::setRotationalPressure, nb::arg("on"),
          "PM I ablation (Guy-Fogelson 2005): False drops the rotational -mu*div(u*) term from "
          "the incremental pressure accumulation (constant-mu path only). Default True = shipped "
          "rotational (Timmermans) update. Also: set_collocated_scheme accepts \"gauge-2a\" -- "
          "the experimental gradient-2a one-sided branch of the gauge-exact gradient.")
      .def(
          "set_rotational_filter", &S::setRotationalFilter, nb::arg("on"), nb::arg("eps") = 0.05,
          "Experimental filtered rotational update: smooth div(u*) (mask-aware axis-wise 1-2-1, "
          "one-sided into the fluid at walls) before accumulating -mu*div into P. Keeps the O(1) "
          "pressure-relaxation gain, removes the checkerboard feedback channel.")
      .def(
          "set_rotational_weight", &S::setRotationalWeight, nb::arg("w"),
          "Under-relax the rotational term: P += ct*phi - w*mu*div(u*). 1 = shipped, 0 = PM I; "
          "small w raises the boundary-mode stability threshold ~1/w at ~1/w slower smooth-mode "
          "pressure relaxation. phi=0 stays the unique fixed point for any w>0 at every dt.")
      .def(
          "set_rotational_wall_weight", &S::setRotationalWallWeight, nb::arg("w0"),
          "Wall-banded rotational blend: at fluid cells with a solid axis-neighbour use "
          "P += (rho/dt + w0*mu/dx^2)*phi - (1-w0)*mu*div(u*); bulk keeps the full rotational "
          "update. Stabilizes the boundary rows without slowing bulk pressure relaxation, and "
          "keeps them relaxing at dt->infinity. 0 (default) = off.")
      .def("set_aperture_order", &S::setApertureOrder, nb::arg("order"),
           "Face-aperture estimator: 1 = one-sample linear model (default, byte-identical), 2 = "
           "marching-squares (5 samples/face, O(h^2); removes the convexity bias measured at "
           "+0.59/+0.27% bed permeability at R=8/12 -- see doc/collocated_paper_plan.md row 51). "
           "For analytic geometry, exact apertures via set_openness_override are better still. "
           "Call before set_solid.")
      .def("set_fluid_only_constraint", &S::setFluidOnlyConstraint, nb::arg("mode"),
           "Fluid-only pressure constraint (route 2b, call before set_solid; collocated "
           "experiment). 1 = Design A (openness filter), 2 = Design B (SPD Kron star "
           "elimination), 0 = off.")
      .def("set_fv_relax", &S::setFvRelax, nb::arg("w"),
           "Mode-4 FV wall-flux defect-correction under-relaxation (1=full; <1 damps the stiff "
           "explicit-lagged wall term). Steady state is independent of w.")
      .def("set_velocity_streams", &S::setVelocityStreams, nb::arg("on"),
           "Toggle overlapped per-component velocity solves.")
      .def("set_implicit_advection", &S::setImplicitAdvection, nb::arg("on"),
           "Use implicit-FOU advection with deferred-correction TVD.")
      .def("set_outer_iterations", &S::setOuterIterations, nb::arg("n"),
           "Set the number of Picard/outer iterations per step.")
      .def("set_outer_tolerance", &S::setOuterTolerance, nb::arg("tol"),
           "Set the outer (Picard) convergence tolerance.")
      .def("last_outer_iterations", &S::lastOuterIterations,
           "Return the outer-iteration count from the last step().")
      .def(
          "set_velocity_solver_params",
          [](S& s, int iters, double rtol, int min_iters) {
            s.setVelocityIterations(iters);
            s.setVelocityTolerance(rtol, min_iters);
          },
          nb::arg("iters"), nb::arg("rtol") = 0.0, nb::arg("min_iters") = 2,
          "Momentum smoother control: `iters` RB-GS sweeps per component, or with rtol > 0 a "
          "TOLERANCE STOP — end the loop once the sweep's max velocity increment has contracted "
          "to rtol of the first sweep's (iters becomes the cap, min_iters the floor). Easy "
          "regimes (small nu*dt/dx^2) exit after ~3-5 sweeps; stiff regimes run to the cap "
          "unchanged. rtol = 0 (default) is the legacy fixed count, byte-identical.")
      .def("set_deferred_correction", &S::setDeferredCorrection, nb::arg("on"),
           "Deferred-correction advection: True (default) = 2nd order (implicit FOU + explicit "
           "high-order correction, the high-order scheme being SOU by default or Koren TVD via "
           "set_advection_scheme); False = pure implicit FOU (1st-order upwind, more dissipative "
           "but "
           "unconditionally stable at sharp shear layers).")
      .def("set_backflow_stabilization", &S::setBackflowStab, nb::arg("beta"),
           "Outflow backflow-stabilization coefficient (Bazilevs 2009 / Esmaily-Moghadam 2011): "
           "beta "
           "in [0,1] scales the dissipative outflow term that prevents backflow divergence when "
           "flow "
           "reverses at the outlet (e.g. a separated wake / BFS recirculation). Default 0.2; 0 = "
           "off. "
           "Inert where the outlet is purely outgoing.")
      .def("set_pressure_solver_params", &S::setPressureIterations, nb::arg("iters"),
           "Set the pressure smoother iteration count.")
      .def(
          "set_pressure_multigrid", [](S& s, bool, int levels) { s.setPressureLevels(levels); },
          nb::arg("on"), nb::arg("levels") = 4,
          "Set the pressure multigrid depth (levels=1 => pure RB-GS, no coarse grid).")
      .def("set_pressure_chebyshev", &S::setPressureChebyshev, nb::arg("on"),
           nb::arg("max_iter") = 120, nb::arg("rtol") = 1e-9,
           "Use the communication-light Chebyshev pressure accelerator (exclusive with PCG).")
      .def(
          "set_pressure_mean_removal",
          [](S& s, const std::string& scope) {
            if (scope != "all" && scope != "fine")
              throw std::runtime_error("set_pressure_mean_removal: scope must be 'all' or 'fine'");
            s.setPressureMeanRemoval(scope == "all");
          },
          nb::arg("scope"),
          "Nullspace (mean) removal scope in the pressure solve: 'fine' (DEFAULT — only the "
          "projections the Krylov iteration needs: rhs/residual, the fine-level V-cycle exit, the "
          "final iterate; ~3x fewer global-reduction latency hits per iteration, the measured "
          "winner of the multi-GPU ablation) or 'all' (legacy — every V-cycle level + after every "
          "matvec). Iteration counts are identical (A preserves mean-freeness); results equal "
          "within solver tolerance, not bit-identical.")
      .def(
          "set_pressure_bottom",
          [](S& s, const std::string& m) {
            if (m == "auto") s.setPressureBottomMode(-1);
            else if (m == "smoother") s.setPressureBottomMode(0);
            else if (m == "agglomerated") s.setPressureBottomMode(1);
            else throw std::runtime_error("set_pressure_bottom: 'auto' | 'smoother' | 'agglomerated'");
          },
          nb::arg("mode"),
          "Coarse-level (bottom) solve of the pressure multigrid. A V-cycle converges at a "
          "domain-independent rate only if its COARSEST level is effectively solved, and a geometric "
          "hierarchy cannot always get small enough: an axis stops coarsening once it turns odd, and "
          "under MPI once ANY rank's block turns odd -- so at fixed cells/rank the coarsest GLOBAL "
          "grid grows with the rank count and the bottom is progressively under-solved. "
          "'auto' (DEFAULT) agglomerates the coarsest level onto a global operator and solves it "
          "exactly whenever it exceeds PECLET_FLOW_AGGLOM_EXTENT (4) cells on any axis, and uses "
          "the cheap smoothed bottom otherwise (byte-identical to 'smoother' then). 'smoother' = "
          "never agglomerate (legacy). 'agglomerated' = always. Measured on one GPU (2048x64x64 "
          "channel): a smoothed bottom needs 13.5 pressure iterations/step at 4 levels and 6.0 at "
          "6, against 4.4 at full geometric depth; agglomerated it is 4.0 at BOTH depths, and "
          "faster in wall-clock than the deep hierarchy. Works on the cut-cell IBM and "
          "ghost-projection paths (per-fluid-component null-space projection, 2026-08-13).")
      .def("set_pressure_graph_amg", &S::setPressureGraphAmg, nb::arg("on"),
           "Solve the pressure MG's coarsest level with an agglomerated mesh-agnostic algebraic "
           "multigrid (core GraphAMG), decomposition-agnostic: with levels=1 this gives a "
           "mesh-independent pressure solve that works under a WEIGHTED ORB (where the geometric "
           "coarse levels can't cleanly coarsen). Applied at the next set_solid.")
      .def("set_pressure_pcg", &S::setPressurePcg, nb::arg("on"), nb::arg("max_iter") = 200,
           nb::arg("rtol") = 1e-8,
           "Use the MG-PCG pressure accelerator (single-GPU default; exclusive with Chebyshev).")
      .def("set_exact_crossings", &S::setExactCrossings, nb::arg("t"),
           "Analytic-SDF capability: exact wall-crossing fractions overriding the "
           "linear-interpolated theta in the momentum cut-cell overlay AND the ghost-projection "
           "closures. Flat array of 9*nx*ny*nz values, blocks [(c*3+k)]: component c's staggered "
           "point at inner cell i toward its +k neighbour; NaN = no crossing. Call BEFORE "
           "set_solid; empty list clears. Single-rank, staggered momentum placement.")
      .def("set_openness_override", &S::setOpennessOverride, nb::arg("ox"), nb::arg("oy"),
           nb::arg("oz"),
           "Analytic-SDF capability: exact face-aperture (openness) fields for the cut-cell "
           "projection, overriding the sampled-SDF openness. Inner nx*ny*nz arrays, x-fastest; "
           "ox[i] = fluid fraction of the -x face of cell i. Call BEFORE set_solid; empty ox "
           "clears. Single-rank.")
      .def("set_ghost_projection", &S::setGhostProjection, nb::arg("on"),
           nb::arg("matrix_order") = 2, nb::arg("rhs_order") = 2,
           "QUARANTINED 2026-08-18 (verification only, unsupported): superseded by the gauge-exact collocated scheme, which matches its accuracy at 5-6x lower cost. Kept as the independent second discretization behind the cross-IBM physics gate. Enabling it on the collocated grid silently selects the plain face map, since it owns the operators the gauge-exact scheme replaces. EXPERIMENTAL directional ghost-cell projection (second staggered IBM): point-based FD "
           "divergence with wall-anchored directional closures instead of the openness-weighted "
           "cut-cell projection; solved by MG-preconditioned BiCGStab. Call BEFORE set_solid. "
           "Closure orders (1=linear, 2=quadratic): (matrix_order, rhs_order) = (2,2) full "
           "quadratic 13-point matrix; (1,1) linear 7-point; (1,2) mixed/deferred — 2nd-order "
           "steady constraint on a 7-point matrix. Collocated: the same closures/matrix on the "
           "face-averaged field, plus a directional (one-sided 2nd-order) cell gradient for the "
           "-grad(P) predictor and cell correction; requires face_interp 0. v1: single-rank, "
           "periodic + IBM, stationary walls; incompatible with "
           "porous/variable-rho/domain-BC/Chebyshev.")
      .def("set_velocity_multigrid", &S::setVelocityMultigrid, nb::arg("on"), nb::arg("levels") = 4,
           nb::arg("vcycles") = 8,
           "Enable velocity (momentum) multigrid for the implicit diffusion solve.")
      .def("last_pressure_iterations", &S::lastPressureIterations,
           "Return the pressure-solver iteration count from the last step().")
      .def(
          "last_step_timers",
          [](S& s) {
            nb::dict d;
            d["step"] = s.lastStepSeconds();
            d["predictor"] = s.lastPredictorSeconds();
            d["momentum"] = s.lastMomentumSeconds();
            d["projection"] = s.lastProjectionSeconds();
            d["pressure_allreduce"] = s.lastPressureAllreduceSeconds();
            d["pressure_allreduce_count"] = s.lastPressureAllreduceCount();
            d["momentum_sweeps"] = s.lastMomentumSweeps();
            return d;
          },
          "Per-phase wall times (seconds, THIS rank, device-fenced) of the last step(): 'predictor' "
          "(ghost fills + RHS/advection/stencil builds), 'momentum' (implicit-diffusion solves), "
          "'projection' (cut-cell pressure projection), 'step' (whole step). "
          "'pressure_allreduce'/'pressure_allreduce_count' = time in / number of global reductions "
          "(MPI_Allreduce) inside the pressure solve — the latency-bound term of the distributed "
          "solve (0 on a single rank).")
      .def("set_domain_bc", &S::setDomainBc, nb::arg("face"), nb::arg("type"), nb::arg("vx") = 0.0,
           nb::arg("vy") = 0.0, nb::arg("vz") = 0.0,
           "Set a per-face domain BC (face 0..5 = -x,+x,-y,+y,-z,+z; type 0 periodic/1 wall/2 "
           "inflow/3 outflow).")
      .def(
          "set_domain_bc_profile",
          [](S& s, int face, nb::ndarray<double, nb::c_contig> prof) {
            if (prof.ndim() != 3 || prof.shape(2) != 3)
              throw std::runtime_error("profile must be (Nb,Nc,3)");
            const int nb_ = (int)prof.shape(0), nc = (int)prof.shape(1);
            s.setDomainBcProfile(
                face, peclet::core::python::ndarray_to_vector<double>(nb::ndarray<>(prof)), nb_,
                nc);
          },
          nb::arg("face"), nb::arg("profile"),
          "Prescribe a per-position inlet velocity profile (Nb,Nc,3) over a face (sets it to "
          "inflow).")
      .def(
          "set_pressure_geometry",
          [](S& s, nb::ndarray<double, nb::f_contig> sdf) { s.setPressureGeometry(grid_in(sdf)); },
          nb::arg("sdf"),
          "Set an all-fluid SDF for the cut-cell pressure operator without an immersed solid (the "
          "channel/BFS domain-BC path). For a no-slip immersed BODY in an inflow/outflow domain, "
          "call "
          "set_solid(sdf, cutcell_pressure=True) instead -- do NOT also call this (a second "
          "geometry "
          "setter overwrites the SDF and wipes the solid).")
      .def(
          "set_solid",
          [](S& s, nb::ndarray<double, nb::f_contig> sdf, bool cutcell_pressure,
             const std::string& /*pressure_coarse*/) {
            s.setSolid(grid_in(sdf), cutcell_pressure);
          },
          nb::arg("sdf"), nb::arg("cutcell_pressure") = false, nb::arg("pressure_coarse") = "const",
          "Set the solid SDF as a Fortran-order (nx,ny,nz) float64 array (negative inside the "
          "solid, positive in fluid). cutcell_pressure=True enables the open-face-weighted "
          "cut-cell "
          "pressure operator (proper no-slip); it composes with domain BCs, so this is the single "
          "call for a no-slip immersed body in an inflow/outflow domain.")
      .def(
          "set_state",
          [](S& s, nb::ndarray<double, nb::f_contig> u, nb::ndarray<double, nb::f_contig> v,
             nb::ndarray<double, nb::f_contig> w) {
            s.uploadVelocity(grid_in(u), grid_in(v), grid_in(w));
          },
          nb::arg("u"), nb::arg("v"), nb::arg("w"),
          "Upload an initial velocity field (u,v,w each a Fortran-order (nx,ny,nz) float64 array).")
      .def("step", &S::step,
           "Advance the solver one time step (semi-implicit: diffusion + projection).")
      .def(
          "get_u", [](S& s) { return field_out(s, s.getVelocity(0)); },
          "Return the x-velocity component as a Fortran-order (nx,ny,nz) float64 array (index "
          "[x,y,z]).")
      .def(
          "get_v", [](S& s) { return field_out(s, s.getVelocity(1)); },
          "Return the y-velocity component as a Fortran-order (nx,ny,nz) float64 array (index "
          "[x,y,z]).")
      .def(
          "get_w", [](S& s) { return field_out(s, s.getVelocity(2)); },
          "Return the z-velocity component as a Fortran-order (nx,ny,nz) float64 array (index "
          "[x,y,z]).")
      .def(
          "get_p", [](S& s) { return field_out(s, s.getPressure()); },
          "Return the physical pressure as a Fortran-order (nx,ny,nz) float64 array (index "
          "[x,y,z]).")
      .def(
          "get_ox", [](S& s) { return field_out(s, s.getOpenness(0)); },
          "TEMP: -x face openness (fluid area fraction) per inner cell, (nx,ny,nz).")
      .def(
          "get_oy", [](S& s) { return field_out(s, s.getOpenness(1)); },
          "TEMP: -y face openness per inner cell, (nx,ny,nz).")
      .def(
          "get_oz", [](S& s) { return field_out(s, s.getOpenness(2)); },
          "TEMP: -z face openness per inner cell, (nx,ny,nz).")
      .def(
          "get_ox_proj", [](S& s) { return field_out(s, s.getOpennessProj(0)); },
          "-x face openness whose fluxes the projection CONSERVES (binary/COUPLED under "
          "set_ghost_projection, geometric cut-cell otherwise). Use for flux bookkeeping "
          "(peclet.pnm extract_network_flow).")
      .def(
          "get_oy_proj", [](S& s) { return field_out(s, s.getOpennessProj(1)); },
          "-y face openness the projection conserves (see get_ox_proj).")
      .def(
          "get_oz_proj", [](S& s) { return field_out(s, s.getOpennessProj(2)); },
          "-z face openness the projection conserves (see get_ox_proj).")
      .def(
          "get_uf", [](S& s) { return field_out(s, s.getFaceVelocity(0)); },
          "Return the divergence-free FACE x-velocity (collocated: projected MAC field; staggered: "
          "== get_u).")
      .def(
          "get_vf", [](S& s) { return field_out(s, s.getFaceVelocity(1)); },
          "Return the divergence-free FACE y-velocity (collocated: projected MAC field; staggered: "
          "== get_v).")
      .def(
          "get_wf", [](S& s) { return field_out(s, s.getFaceVelocity(2)); },
          "Return the divergence-free FACE z-velocity (collocated: projected MAC field; staggered: "
          "== get_w).")
      // --- Named field registry (multiphysics field container) ---------------------------------
      .def(
          "add_field", [](S& s, const std::string& name) { s.addField(name); }, nb::arg("name"),
          "Register a new zero-initialised cell-centred field on the grid (for transported scalars "
          "or material properties). Idempotent.")
      .def(
          "has_field", [](S& s, const std::string& name) { return s.hasField(name); },
          nb::arg("name"), "Whether a field of this name is registered.")
      .def(
          "field_names", [](S& s) { return s.fieldNames(); },
          "Names of all registered fields (velocity u/v/w, p, sdf, plus any added), sorted.")
      .def(
          "get_field", [](S& s, const std::string& name) { return field_out(s, s.getField(name)); },
          nb::arg("name"),
          "Return a registered field's inner region as a Fortran-order (nx,ny,nz) float64 array.")
      .def(
          "set_field",
          [](S& s, const std::string& name, nb::ndarray<double, nb::f_contig> a) {
            s.setField(name, grid_in(a));
          },
          nb::arg("name"), nb::arg("array"),
          "Write a Fortran-order (nx,ny,nz) float64 array into a registered field's inner region "
          "(ghosts refilled on the next exchange_field/step).")
      .def(
          "field_view",
          [](S& s, const std::string& name) { return field3d_out(s, s.fieldView(name)); },
          nb::arg("name"),
          "Zero-copy view of a registered field's full padded buffer as a Fortran-order "
          "(nx+2g, ny+2g, nz+2g) array (g = ghost_width); host → NumPy, device → DLPack (CuPy).")
      .def(
          "exchange_field", [](S& s, const std::string& name) { s.exchangeField(name); },
          nb::arg("name"),
          "Fill a registered field's ghost cells (cross-rank + periodic under MPI; periodic "
          "single-rank).")
      .def(
          "exchange_field_add", [](S& s, const std::string& name) { s.exchangeFieldAdd(name); },
          nb::arg("name"),
          "Add-reduce halo: fold ghost-layer deposits back onto their owner (cross-rank + "
          "periodic). "
          "The particle->grid deposition primitive for MPI CFD-DEM; single-rank non-periodic "
          "no-op.")
      // --- Scalar transport (advection-diffusion) ----------------------------------------------
      .def(
          "add_scalar",
          [](S& s, const std::string& name, double diffusivity, int scheme, int iters) {
            s.addScalar(name, diffusivity, scheme, iters);
          },
          nb::arg("name"), nb::arg("diffusivity") = 0.0, nb::arg("scheme") = 1,
          nb::arg("iters") = 50,
          "Register a transported scalar (temperature/concentration/…): constant diffusivity (grid "
          "units), advection scheme 0=FOU/1=Koren TVD/2=SOU, and RB-GS diffusion sweeps. The "
          "scalar "
          "is a registered field (get_field/set_field/field_view). Requires geometry "
          "(set_solid/set_pressure_geometry) for the openness-weighted operators.")
      .def(
          "set_scalar_bc",
          [](S& s, const std::string& name, int face, int type, double value) {
            s.setScalarBc(name, face, type, value);
          },
          nb::arg("name"), nb::arg("face"), nb::arg("type"), nb::arg("value") = 0.0,
          "Scalar boundary condition on a domain face (0..5 = -x,+x,-y,+y,-z,+z): type 0 periodic, "
          "1 Neumann zero-flux (adiabatic), 2 Dirichlet value. Single-rank.")
      .def(
          "has_scalar", [](S& s, const std::string& name) { return s.hasScalar(name); },
          nb::arg("name"), "Whether a transported scalar of this name is registered.")
      .def(
          "advance_scalars", [](S& s) { s.advanceScalars(); },
          "Advance all registered scalars one dt with the current velocity (also done by step()).")
      // --- Property closures + Boussinesq body force -------------------------------------------
      .def(
          "set_property_model",
          [](S& s, const std::string& target, const std::string& kind, const std::string& in0,
             const std::vector<double>& params, const std::string& in1) {
            peclet::flow::ClosureKind k;
            if (kind == "linear")
              k = peclet::flow::ClosureKind::LinearMix;
            else if (kind == "boussinesq")
              k = peclet::flow::ClosureKind::BoussinesqForce;
            else if (kind == "arrhenius")
              k = peclet::flow::ClosureKind::ArrheniusMu;
            else
              throw std::runtime_error("set_property_model: unknown kind '" + kind + "'");
            s.setPropertyModel(target, k, in0, in1, params);
          },
          nb::arg("target"), nb::arg("kind"), nb::arg("field"),
          nb::arg("params") = std::vector<double>{}, nb::arg("field2") = std::string{},
          "Register a device closure writing a property/body-force field from input field(s). "
          "target: a registered field (a property 'mu'/'rho'/… or a body-force component "
          "'force_x'/'force_y'/'force_z'). kind: 'linear' (params [p0,p1,p2]: "
          "p0+p1*field+p2*field2), "
          "'boussinesq' (params [rho0,g,beta,T0]: rho0*g*beta*(field-T0) buoyancy), 'arrhenius' "
          "(params [mu_ref,B,Tref]: mu_ref*exp(B*(1/field-1/Tref))). Applied at the top of step().")
      .def(
          "set_property_table",
          [](S& s, const std::string& target, const std::string& field,
             const std::vector<double>& x,
             const std::vector<double>& y) { s.setPropertyTable(target, field, x, y); },
          nb::arg("target"), nb::arg("field"), nb::arg("x"), nb::arg("y"),
          "Register a tabulated property: target = piecewise-linear interpolation of (x, y) at the "
          "input field value (x ascending, clamped at the ends).")
      .def(
          "update_properties", [](S& s) { s.updateProperties(); },
          "Apply all registered property/force closures now (also done at the top of step()).")
      .def(
          "enable_cell_force", [](S& s) { s.enableCellForce(); },
          "Allocate + register the per-cell body-force fields force_x/force_y/force_z and route "
          "them "
          "into the momentum RHS, for an external writer (e.g. CFD-DEM drag feedback) to fill "
          "directly via field_view('force_z'). They persist across steps until overwritten.")
      .def(
          "enable_drag", [](S& s) { s.enableDrag(); },
          "Enable implicit (semi-implicit) linear drag for CFD-DEM: allocate the per-cell "
          "'drag_beta' "
          "field (added to the momentum diagonal so a -beta*(u-u_p) source is treated implicitly "
          "-> "
          "unconditionally stable for the stiff beta of a dense bed) plus force_x/y/z (which carry "
          "beta*u_p, the RHS target). Fill 'drag_beta' and 'force_*' via field_view each step.")
      .def(
          "set_property_mode",
          [](S& s, const std::string& mode, bool harmonic) {
            s.setPropertyMode(mode == "variable", harmonic);
          },
          nb::arg("mode") = "variable", nb::arg("harmonic") = false,
          "Enable variable-coefficient momentum (variable viscosity): mode 'variable' binds the "
          "'mu' "
          "field (get/set_field('mu')) into the diffusion operator; 'constant' reverts. harmonic = "
          "harmonic face-viscosity mean (continuous shear stress across a jump) vs arithmetic. A "
          "closure targeting 'mu' enables this automatically. The incremental-rotational pressure "
          "scheme (large-dt / steady-Stokes) stays active — see set_variable_rotational.")
      .def(
          "set_variable_rotational",
          [](S& s, const std::string& mode, double chi) {
            int m = 0;
            if (mode == "min")
              m = 0;
            else if (mode == "full")
              m = 1;
            else if (mode == "off")
              m = 2;
            else
              throw std::runtime_error("set_variable_rotational: mode must be min/full/off");
            s.setVariableRotational(m, chi);
          },
          nb::arg("mode") = "min", nb::arg("chi") = 1.0,
          "Rotational-pressure term under variable viscosity (the constant-mu Timmermans term "
          "-mu*div(u*) is only valid for homogeneous viscosity — Deteix & Yakoubi 2018). 'min' "
          "(default): constant coefficient chi*mu_min — provably stable at any contrast, exact "
          "fallback to the constant-mu scheme for uniform mu. 'full': pointwise chi*mu(i) — better "
          "pressure consistency at MILD contrast only. 'off': plain incremental (no rotational "
          "term). All modes keep the incremental predictor (large-dt / steady-Stokes capability).")
      .def(
          "set_density_mode",
          [](S& s, const std::string& mode) { s.setDensityMode(mode == "variable"); },
          nb::arg("mode") = "variable",
          "Enable variable density (staggered solver only): binds the 'rho' field "
          "(get/set_field('rho'), created seeded with set_rho's value if absent) into the momentum "
          "time term, the advection weight, the per-cell body force (face-interpolated), and the "
          "pressure projection (face coefficient openness*rho0/rho_f with the matching 1/rho_f "
          "velocity correction; rho0 = set_rho's value, so a uniform field reduces exactly to the "
          "constant solver). A closure targeting 'rho' (e.g. a linear mixture of a transported "
          "phase fraction) enables this automatically. For gravity, register a closure "
          "force_z = linear(rho, params=[0, -g]).")
      .def(
          "ghost_width", [](S& s) { return s.ghostWidth(); },
          "Ghost-layer width g of the velocity block (field_view returns an (n+2g) buffer).")
      .def(
          "has_cutcell_pressure", [](S& s) { return s.hasCutcellPressure(); },
          "True once the cut-cell pressure operator exists (set_solid or set_pressure_geometry was "
          "called). The porous continuity requires it; the coupling driver auto-sets an all-fluid "
          "geometry when absent.")
      .def(
          "set_porous_continuity", [](S& s, bool on) { s.setPorousContinuity(on); },
          nb::arg("on") = true,
          "Enable the volume-averaged (porous) continuity for unresolved CFD-DEM (staggered only): "
          "the projection enforces d(eps)/dt + div(eps u) = 0 instead of div(u)=0, so the fluid "
          "velocity is NOT solenoidal where the void fraction changes (bubbling/expansion). Binds "
          "the 'eps' field (void fraction from the particle deposition, created seeded to 1 if "
          "absent; the coupling writes it each step BEFORE step()). eps=1 everywhere reduces "
          "exactly "
          "to div(u)=0. Pair with max_porous_residual() for the meaningful convergence check.")
      .def("max_open_divergence", &S::maxOpenDivergence,
           "Return the max cut-cell velocity-flux divergence max|div(open*u)|. With porous "
           "continuity this is NOT ~0 -- it equals -d(eps)/dt (the bed expanding). Use "
           "max_porous_residual() for the continuity residual.")
      .def(
          "set_pressure_underrelax", [](S& s, double w) { s.setPressureUnderRelax(w); },
          nb::arg("omega"),
          "Pressure under-relaxation factor omega_p in (0,1] for the incremental accumulation "
          "(MFIX "
          "§10.1); 1.0 = off (default). <1 damps the incremental predictor overshoot on stiff "
          "porous+drag.")
      .def(
          "set_porous_deps_dt", [](S& s, bool on) { s.setPorousDepsDt(on); }, nb::arg("on"),
          "Include (default True) or drop the d(eps)/dt source in the porous projection RHS. Drop "
          "it "
          "to enforce div(eps u)=0 when the per-cell eps deposit's time-derivative is too jagged "
          "and "
          "destabilizes the eps-weighted pressure solve.")
      .def(
          "set_porous_conservative", [](S& s, bool on) { s.setPorousConservative(on); },
          nb::arg("on"),
          "eps-conservative porous momentum + projection pair (default True): time term "
          "(eps_f rho/dt) u, eps rho-weighted advective form, projection coefficients "
          "open*(eps rho idt)/(eps rho idt+beta) with matching correction. False = the legacy "
          "plain-u pair (A/B only; it lets the projection drag gas with the moving porosity at "
          "zero inertia cost — a spurious late-time energy source in clustering flows).")
      .def("sync_porous_prev", &S::syncPorousPrev,
           "Reseed eps^n = eps^{n+1} (d(eps)/dt=0 this step) — call once after the first "
           "void-fraction "
           "deposition so step 0 has no spurious source.")
      .def("max_porous_residual", &S::maxPorousResidual,
           "Residual of the volume-averaged continuity max|div(open*eps*u) + d(eps)/dt| -- the "
           "quantity the porous projection drives to zero. 0 unless set_porous_continuity(True).")
      .def(
          "get_resolution", [](S& s) { return std::vector<int>{s.nx(), s.ny(), s.nz()}; },
          "Return the LOCAL grid resolution [nx, ny, nz] (this rank's block under MPI).")
      .def(
          "global_resolution",
          [](S& s) {
            auto g = s.globalResolution();
            return std::vector<int>{g[0], g[1], g[2]};
          },
          "Return the GLOBAL grid resolution [gnx, gny, gnz] (== local single-rank). For the "
          "CFD-DEM co-decomposition weight field.")
      .def(
          "block_origin",
          [](S& s) {
            auto o = s.blockOrigin();
            return std::vector<int>{o[0], o[1], o[2]};
          },
          "This rank's inner-block origin in GLOBAL cells ([0,0,0] single-rank). Shift the "
          "coupling "
          "deposit origin by this so particles in global coordinates land in the local block.")
      .def(
          "get_spacing", [](S&) { return std::vector<double>{1.0, 1.0, 1.0}; },
          "Return the grid spacing [dx, dy, dz] (always unit on this grid).")
#ifdef PECLET_FLOW_MPI
      // Distributed path (built with -DPECLET_FLOW_MPI): construct the Solver with this rank's
      // LOCAL block dims (see the module-level mpi_block()), then init_mpi with the GLOBAL grid
      // dims. step() then does the g=2 velocity-block halo exchange + the distributed cut-cell
      // pressure MG. Bit-exact to single-rank.
      .def(
          "init_mpi",
          [](S& s, int gnx, int gny, int gnz) {
            ensure_mpi_init();
            s.initMpi(gnx, gny, gnz, MPI_COMM_WORLD);
          },
          nb::arg("gnx"), nb::arg("gny"), nb::arg("gnz"),
          "Wire the multi-rank step: pass the GLOBAL grid dims (gnx,gny,gnz). The Solver must have "
          "been "
          "constructed with this rank's LOCAL block dims (from mpi_block). MPI_Init is called if "
          "needed.")
      .def(
          "rebalance_by_weights",
          [](S& s, const std::vector<double>& w) { s.rebalanceByWeights(w); }, nb::arg("weights"),
          "Dynamic load balancing: redistribute the solver's state onto the weighted ORB of "
          "per-cell "
          "weights (global x-fastest, gnx*gny*gnz). Pass fluid work + gamma*particle_count and the "
          "coupled dem migrates onto the SAME partition from the same array. State-preserving "
          "(bit-exact at np=1, reduction floor at np>1).")
      .def(
          "rank",
          [](S&) {
            ensure_mpi_init();
            int r = 0;
            MPI_Comm_rank(MPI_COMM_WORLD, &r);
            return r;
          },
          "This rank's index in MPI_COMM_WORLD.")
      .def(
          "size",
          [](S&) {
            ensure_mpi_init();
            int n = 1;
            MPI_Comm_size(MPI_COMM_WORLD, &n);
            return n;
          },
          "The number of ranks in MPI_COMM_WORLD.")
#else
      .def(
          "rank", [](S&) { return 0; },
          "MPI rank (always 0 in the single-rank Python module; the multi-rank path is the "
          "tests/kokkos_mpi suite).")
      .def(
          "size", [](S&) { return 1; }, "MPI size (1 in the single-rank Python module).")
#endif
      .def(
          "bcast_from_root", [](S&, nb::object v) { return v; }, nb::arg("value"),
          "Broadcast a value from rank 0 (identity in the single-rank module; mirrors the MPI "
          "API).");
}
 
NB_MODULE(_flow, m) {
  m.attr("__doc__") =
      "flow — Kokkos cut-cell IBM incompressible Navier-Stokes solver for porous media.\n\n"
      "Two solver classes share an identical API (only the velocity-unknown placement differs):\n"
      "  Solver          — staggered MAC grid (THE flow solver; permeability/drag accuracy "
      "default)\n"
      "  SolverColocated — collocated / cell-centered velocities (ABC approximate projection)\n\n"
      "Conventions: physical units throughout (density rho, viscosity mu, physical pressure p); "
      "SDFs\n"
      "are negative inside the solid; fields are Fortran-order (nx,ny,nz) float64 (x-fastest). "
      "This is\n"
      "the single-rank module — the multi-rank MPI path is exercised by the tests/kokkos_mpi "
      "suite.\n\n"
      "Kokkos is initialized at import and finalized via a Python atexit hook. Release every "
      "Solver "
      "before interpreter exit (it goes out of scope, or `del s; gc.collect()`) so no Kokkos View "
      "outlives finalize.";
  if (!Kokkos::is_initialized())
    Kokkos::initialize();
  // Register Kokkos::finalize via Python atexit. This is REQUIRED on CUDA: without it, Kokkos's
  // internal device state is torn down by static destructors AFTER the CUDA runtime unloads,
  // aborting with cudaErrorCudartUnloading at every exit. atexit runs the hook while the driver is
  // still up. (Returned fields are backed by host std::vectors, not device Views, so they never
  // block finalize; a live Solver still holding Views at exit must be released first — hence the
  // docstring note.)
  nb::module_::import_("atexit").attr("register")(nb::cpp_function([]() {
    if (Kokkos::is_initialized() && !Kokkos::is_finalized())
      Kokkos::finalize();
  }));
  // The active Kokkos backend ("OpenMP", "Cuda", "HIP"), chosen by the build's install prefix.
  m.attr("execution_space") = nb::str(Kokkos::DefaultExecutionSpace::name());
 
  // Staggered MAC grid (THE flow solver) + the collocated/cell-centered variant. Same Python API.
  bind_solver<peclet::flow::Staggered>(m, "Solver");
  bind_solver<peclet::flow::Colocated>(m, "SolverColocated");
 
#ifdef PECLET_FLOW_MPI
  // Module-level: this rank's ORB block of the global (gnx,gny,gnz) grid, matching the
  // deterministic BlockDecomposer the Solver's initMpi re-derives internally (and the C++
  // tests/kokkos_mpi template). Returns (origin=[ox,oy,oz], size=[lnx,lny,lnz]); slice the global
  // SDF with these to build the local block, then Solver(*size) + init_mpi(gnx,gny,gnz). MPI_Init
  // is called if needed.
  m.def(
      "mpi_block",
      [](int gnx, int gny, int gnz) {
        ensure_mpi_init();
        int rank = 0, size = 1;
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);
        MPI_Comm_size(MPI_COMM_WORLD, &size);
        // Derive it through the SAME factory Solver::initMpi uses, so this local block size matches
        // the solver's dec_ under either decomposition mode (see set_decomposition_levels).
        auto dec = peclet::flow::CutcellMG::decomposition(static_cast<std::size_t>(size), gnx, gny,
                                                          gnz);
        auto blk = dec.block(static_cast<std::size_t>(rank));
        std::vector<int> origin{(int)blk.origin[0], (int)blk.origin[1], (int)blk.origin[2]};
        std::vector<int> bsize{(int)blk.size[0], (int)blk.size[1], (int)blk.size[2]};
        return std::make_pair(origin, bsize);
      },
      nb::arg("gnx"), nb::arg("gny"), nb::arg("gnz"),
      "Return this MPI rank's ORB block of the global (gnx,gny,gnz) grid as (origin, size), each a "
      "length-3 list [x,y,z]. Use it to slice the global SDF into this rank's local block for a "
      "distributed Solver (see Solver.init_mpi). MPI_Init is called if needed.");
 
  m.def(
      "set_decomposition_levels",
      [](int levels) { peclet::flow::CutcellMG::setDecompositionLevels(levels); },
      nb::arg("levels"),
      "Choose how the shared MPI decomposition is built. 0 (default) = the aligned ORB: split "
      "positions on the FINE grid are snapped to a power of two (capped at 16, i.e. 5 nested "
      "multigrid levels). levels >= 2 = COARSE-FIRST: build the ORB on the grid coarsened "
      "levels-1 times and refine the partition upward, so every block is a multiple of the "
      "coarsening factor BY CONSTRUCTION and the hierarchy nests for the full requested depth. "
      "Coarse-first also balances better: the aligned ORB picks a split and then rounds it (a "
      "balanced 96|96 can round to 128|64), whereas on the coarse grid one cell IS the quantum, so "
      "the rank count no longer has to be a power of two -- what matters is that the COARSE grid "
      "divides among the ranks. Backs off automatically if the coarse grid would have fewer cells "
      "than ranks. CALL BEFORE mpi_block() AND Solver.init_mpi(): both derive the same partition "
      "from this setting and must agree. Env override: PECLET_FLOW_DECOMP_LEVELS.");
  m.def("decomposition_levels", [] { return peclet::flow::CutcellMG::decompositionLevels(); },
        "Current decomposition mode (0 = aligned ORB, >= 2 = coarse-first with that depth).");
 
  m.attr("has_mpi") = true;
#else
  m.attr("has_mpi") = false;
#endif
}