amr_bindings.cpp

Go to the documentation of this file.

// amr — Python surface of peclet::amr (the `peclet.amr` package).
//
// A nanobind module exposing the adaptive-mesh-refinement package: the per-block BlockOctree
// (serial) and the MPI DistributedOctree (ORB over root cells), so an mpi4py driver can build a
// graded octree, refine it to a signed-distance surface, read leaf geometry + a per-leaf field as
// numpy, load-rebalance the distributed octree, gather face-neighbour values, and export VTU; the
// geometric-multigrid Poisson solver; and the device (Kokkos) collocated Navier–Stokes `Flow`,
// whose developer tier is `Flow.diagnostics` (FlowDiagnostics, suite QUALITY_PLAN D2).
//
// Conventions (see ../docs/CONVENTIONS.md):
//   * 3D only here (Bits=21, codes in a 64-bit int). x-fastest where a linear index appears.
//   * Leaf arrays are in the octree's Z-order slot order; every per-leaf numpy array (centers,
//     sizes, levels, codes, fields) is indexed by that same slot, length num_leaves.
//   * Geometry: a leaf's world centre = origin + spacing * fine_centre (per axis); a leaf at
//     refinement `level` is spacing * 2**level wide. Root cells sit at level=lmax; level decreases
//     toward 0 (finest). `cells` is always the FINEST grid (root cells * 2**lmax), as in flow.
//
// Per-leaf arrays are returned via the shared peclet::core::python::vector_to_ndarray
// (capsule-backed, no extra copy); the Flow path returns host fields the same way. The module links
// MPI but never calls MPI_Init/Finalize: a distributed driver imports mpi4py.MPI first, and the
// distributed classes run on MPI_COMM_WORLD.
//
// Build: see CMakeLists.txt (the `amr` target -> peclet.amr._amr, re-exported by
// packaging/amr_init.py).
#include <mpi.h>
#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
#include <nanobind/stl/array.h>
#include <nanobind/stl/function.h>
#include <nanobind/stl/optional.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/tuple.h>
#include <nanobind/stl/variant.h>
#include <nanobind/stl/vector.h>
 
#include <array>
#include <cstdint>
#include <functional>
#include <Kokkos_Core.hpp>
#include <optional>
#include <set>
#include <string>
#include <vector>
 
#include "peclet/amr/adapt.hpp"
#include "peclet/amr/block_octree.hpp"
#include "peclet/amr/distributed_adapt.hpp"
#include "peclet/amr/distributed_octree.hpp"
#include "peclet/amr/flow.hpp"         // the canonical (device) AmrFlow exposed to Python
#include "peclet/amr/flow_oracle.hpp"  // oracle::AmrFlow (dev-only reference; not exposed)
#include "peclet/amr/indicators.hpp"
#include "peclet/amr/leaf_field.hpp"
#include "peclet/amr/mg_predict.hpp"  // predictPressureLadder (docs/amr_mg_depth.md §6.7)
#include "peclet/amr/poisson.hpp"
#include "peclet/amr/refine.hpp"
#include "peclet/amr/vtu_io.hpp"
#include "peclet/core/common/types.hpp"
#include "peclet/core/geom/scene_query.hpp"  // SphereBedQuery (Layer 2-for-core)
#include "peclet/core/geom/sdf.hpp"
#include "peclet/core/python/kokkos_teardown.hpp"
#include "peclet/core/python/ndarray_interop.hpp"
 
namespace nb = nanobind;
using namespace peclet::core;
 
// The wrapper classes live in a NAMED namespace on purpose: they have virtual methods (Releasable),
// and with an anonymous namespace hipcc's host/device split leaves their vtables unemitted, so the
// HIP link fails with `ld.lld: undefined hidden symbol: vtable for (anonymous namespace)::Octree`
// (and Poisson / Flow / DistributedOctree). OpenMP/CUDA builds are indifferent.
namespace peclet::amr::pybind {
 
using BO = amr::BlockOctree<3>;  // 3D, Bits=21 (default) — codes fit a uint64
using DO = amr::DistributedOctree<3>;
using Code = BO::Code;
using Coord = BO::Coord;
using DArray = nb::ndarray<double, nb::c_contig>;
 
// ---- per-leaf array helpers (capsule-backed; no extra copy; Z-order slot order) ----------------
 
template <class T>
nb::ndarray<nb::numpy, T> vec1(std::vector<T> v) {
  const std::size_t n = v.size();
  return peclet::core::python::vector_to_ndarray<T>(std::move(v), {n}, {1});
}
template <class T>
nb::ndarray<nb::numpy, T> vec2(std::vector<T> v, std::size_t cols) {
  const std::size_t n = v.size() / cols;
  return peclet::core::python::vector_to_ndarray<T>(std::move(v), {n, cols},
                                                    {static_cast<std::int64_t>(cols), 1});
}
 
// World centre of every leaf -> (N,3) float64.
nb::ndarray<nb::numpy, double> leafCenters(const BO& t, const amr::AmrGeometry<3>& geo) {
  const Index n = t.numLeaves();
  std::vector<double> d(static_cast<std::size_t>(n) * 3);
  for (Index i = 0; i < n; ++i) {
    Vec<3> c = geo.center(t.bounds(i));
    d[i * 3 + 0] = c[0];
    d[i * 3 + 1] = c[1];
    d[i * 3 + 2] = c[2];
  }
  return vec2<double>(std::move(d), 3);
}
 
// World width of every leaf along `axis` (spacing[axis] * 2**level) -> (N,) float64.
// Phase 3: a leaf is a BOX, so a single width exists only on a cubic octree; `axis` selects one.
nb::ndarray<nb::numpy, double> leafSizes(const BO& t, const amr::AmrGeometry<3>& geo, int axis) {
  if (axis < 0 || axis > 2)
    throw std::runtime_error("sizes(axis): axis must be 0, 1 or 2");
  const Index n = t.numLeaves();
  std::vector<double> d(static_cast<std::size_t>(n));
  for (Index i = 0; i < n; ++i)
    d[i] = geo.leafSize(t.level(i), axis);
  return vec1<double>(std::move(d));
}
 
// Refinement level of every leaf -> (N,) int32 (lmax at the root, 0 finest).
nb::ndarray<nb::numpy, std::int32_t> leafLevels(const BO& t) {
  const Index n = t.numLeaves();
  std::vector<std::int32_t> d(static_cast<std::size_t>(n));
  for (Index i = 0; i < n; ++i)
    d[i] = static_cast<std::int32_t>(t.level(i));
  return vec1<std::int32_t>(std::move(d));
}
 
// Block-local Morton origin code of every leaf -> (N,) uint64.
nb::ndarray<nb::numpy, std::uint64_t> leafCodes(const BO& t) {
  const Index n = t.numLeaves();
  std::vector<std::uint64_t> d(static_cast<std::size_t>(n));
  for (Index i = 0; i < n; ++i)
    d[i] = static_cast<std::uint64_t>(t.code(i));
  return vec1<std::uint64_t>(std::move(d));
}
 
// A contiguous vector<double> as a (N,) float64 numpy array.
nb::ndarray<nb::numpy, double> vecToArray(std::vector<double> v) {
  return vec1<double>(std::move(v));
}
 
// Validate a per-leaf field numpy array and view it as a contiguous vector<double>.
std::vector<double> asField(const BO& t, const DArray& field, const char* who) {
  if (static_cast<Index>(field.shape(0)) != t.numLeaves())
    throw std::runtime_error(std::string(who) + ": field length != num_leaves");
  const double* f = field.data();
  return std::vector<double>(f, f + static_cast<std::size_t>(t.numLeaves()));
}
 
// ---- teardown registry -------------------------------------------------------------------------
// The suite-wide pattern (peclet/core/python/kokkos_teardown.hpp): Kokkos::finalize() MUST run from
// the atexit hook (else cudaErrorCudartUnloading on CUDA), but any wrapper that still holds Views
// at that point -- an Octree/Flow at script or notebook scope, the normal case -- is destroyed
// after finalize and Kokkos::abort()s the process. Every wrapper below is a Releasable (registered
// on construction); the hook drops their Views (release()) BEFORE finalize.
using peclet::core::python::Releasable;
 
// ---- serial single-block octree ----------------------------------------------------------------
 
// Derive the finest cell width h0 from a PHYSICAL extent and the finest cell counts
// (suite/docs/PHYSICAL_UNITS_PLAN.md, work order U6): the caller states the box, the octree
// derives its own spacing, and nobody writes h. Phase 1 is isotropic — the octree's leaves are
// cubes by construction (a level-l leaf is h0*2**l wide on every axis), so an extent that does not
// give one spacing is refused with the numbers rather than silently taking the x one.
// PHASE 3 (suite/docs/PHYSICAL_UNITS_PLAN.md, core/docs/amr_anisotropic.md): the octree's cells
// are BOXES. `extent/cells` is taken per axis and NOTHING is refused — the octree refines by 2 on
// every axis, so a level-l leaf is `h0[d]*2**l` wide on axis `d` and every level inherits the root
// aspect ratio.
inline std::array<double, 3> spacingsFromExtent(const std::array<double, 3>& extent,
                                                const std::array<long, 3>& cells) {
  std::array<double, 3> h{};
  for (int d = 0; d < 3; ++d) {
    if (!(extent[d] > 0.0))
      throw std::runtime_error("extent: every component must be > 0");
    if (cells[d] <= 0)
      throw std::runtime_error("extent: every finest-cell count must be > 0");
    h[d] = extent[d] / (double)cells[d];
  }
  return h;
}
 
// The SCALAR helper keeps its cubic contract: it returns one number, so an extent that does not
// give one spacing is still refused with the three values. `spacings_from_extent` is the per-axis
// form, and the constructors take the anisotropic extent directly.
inline double spacingFromExtent(const std::array<double, 3>& extent,
                                const std::array<long, 3>& cells) {
  const std::array<double, 3> h = spacingsFromExtent(extent, cells);
  for (int d = 1; d < 3; ++d)
    if (std::fabs(h[d] - h[0]) > 1e-12 * std::fabs(h[0])) {
      char msg[320];
      std::snprintf(msg, sizeof msg,
                    "spacing_from_extent returns ONE number, but extent/cells gives "
                    "dx=%.17g dy=%.17g dz=%.17g — the octree's cells are CUBES here. Use "
                    "spacings_from_extent (or pass `extent` to the constructor) for a box mesh.",
                    h[0], h[1], h[2]);
      throw std::runtime_error(msg);
    }
  return h[0];
}
 
// `spacing=` accepts one number (cubic cells) or three (a box mesh).
using SpacingArg = std::variant<double, std::array<double, 3>>;
inline std::array<double, 3> spacingOf(const SpacingArg& sp) {
  if (const double* d = std::get_if<double>(&sp))
    return {*d, *d, *d};
  return std::get<std::array<double, 3>>(sp);
}
 
// The octree is stated on its FINEST grid (`cells`, the same quantity as flow's `cells`); the
// root brick it is built from is cells / 2**lmax, which must divide exactly.
inline std::array<long, 3> rootCellsOf(const std::array<long, 3>& cells, unsigned lmax,
                                       const char* who) {
  const long f = (long)1 << lmax;
  std::array<long, 3> root{};
  for (int d = 0; d < 3; ++d) {
    if (cells[d] <= 0 || cells[d] % f != 0) {
      char msg[256];
      std::snprintf(msg, sizeof msg,
                    "%s: cells=(%ld, %ld, %ld) must be positive multiples of 2**lmax = %ld "
                    "(`cells` is the FINEST grid; the root brick is cells / 2**lmax)",
                    who, cells[0], cells[1], cells[2], f);
      throw std::runtime_error(msg);
    }
    root[d] = cells[d] / f;
  }
  return root;
}
 
// Resolve the constructor's `spacing=` / `extent=` pair into the per-axis finest spacing.
inline std::array<double, 3> resolveSpacing(const std::optional<SpacingArg>& spacing,
                                            const std::optional<std::array<double, 3>>& extent,
                                            const std::array<long, 3>& cells, const char* who) {
  if (spacing && extent)
    throw std::runtime_error(std::string(who) +
                             ": pass EITHER spacing= (the finest cell size) OR extent= (the box "
                             "side lengths; spacing = extent / cells), not both");
  if (extent)
    return spacingsFromExtent(*extent, cells);
  if (spacing)
    return spacingOf(*spacing);
  return {1.0, 1.0, 1.0};
}
 
// A per-block adaptive octree with its world placement (origin + per-axis finest spacing).
// Wraps peclet::amr::BlockOctree<3> + AmrGeometry<3>: build a uniform brick, refine toward a
// surface, query leaves, and read leaf geometry / fields as numpy. The serial / single-rank form;
// for the distributed (MPI) octree use DistributedOctree below.
class Octree : public Releasable {
 public:
  Octree(std::array<long, 3> root, unsigned lmax, std::array<double, 3> origin,
         std::array<double, 3> spacing) {
    t_.init(IVec<3>{root[0], root[1], root[2]}, lmax);
    geo_.origin = {origin[0], origin[1], origin[2]};
    geo_.h0 = {spacing[0], spacing[1], spacing[2]};
  }
  void release() noexcept override { t_ = BO{}; }
 
  // ---- introspection ----
  Index num_leaves() const { return t_.numLeaves(); }
  unsigned lmax() const { return t_.lmax(); }
  std::array<double, 3> origin() const { return {geo_.origin[0], geo_.origin[1], geo_.origin[2]}; }
  std::array<long, 3> cells() const {
    const auto b = t_.brick();
    const long f = (long)1 << t_.lmax();
    return {(long)b[0] * f, (long)b[1] * f, (long)b[2] * f};
  }
  std::array<double, 3> extent() const {
    const auto c = cells();
    return {(double)c[0] * geo_.h0[0], (double)c[1] * geo_.h0[1], (double)c[2] * geo_.h0[2]};
  }
  std::array<double, 3> spacing() const { return {geo_.h0[0], geo_.h0[1], geo_.h0[2]}; }
  bool is_balanced() const { return t_.isBalanced(); }
 
  nb::ndarray<nb::numpy, double> centers() const { return leafCenters(t_, geo_); }
  nb::ndarray<nb::numpy, double> sizes(int axis) const { return leafSizes(t_, geo_, axis); }
  nb::ndarray<nb::numpy, std::int32_t> levels() const { return leafLevels(t_); }
  nb::ndarray<nb::numpy, std::uint64_t> codes() const { return leafCodes(t_); }
 
  // ---- queries ----
  // Index of the leaf containing world point (x,y,z), or -1 if outside the block.
  Index find(std::array<double, 3> x) const {
    std::array<Coord, 3> fine{};
    for (int d = 0; d < 3; ++d) {
      double f = (x[d] - geo_.origin[d]) / geo_.h0[d];
      if (f < 0)
        return -1;
      fine[d] = static_cast<Coord>(f);
    }
    return t_.find(fine);
  }
 
  // ---- refinement ----
  // Refine leaves the sphere's surface passes through (plus a band) down to target_level.
  Index refine_to_sphere(std::array<double, 3> center, double radius, unsigned target_level,
                         double band, bool balance) {
    geom::Sphere s{{center[0], center[1], center[2]}, radius};
    return amr::refineToSdf(
        t_, geo_, [&](const Vec<3>& p) { return s.eval(p); }, target_level, band, balance);
  }
 
  // Refine toward an arbitrary signed-distance field given as a Python callable f(x,y,z)->distance
  // (suite SDF sign: <0 inside solid). Convenient for scripting; the callback runs under the GIL.
  Index refine_to_sdf(std::function<double(double, double, double)> sdf, unsigned target_level,
                      double band, bool balance) {
    return amr::refineToSdf(
        t_, geo_, [&](const Vec<3>& p) { return sdf(p[0], p[1], p[2]); }, target_level, band,
        balance);
  }
 
  // Graded surface refinement: a per-point TARGET LEVEL callable instead of one global target
  // (the mixed-level cut band, plan §7). Pair with Flow.set_ghost_sampled(True).
  Index refine_to_sdf_graded(std::function<double(double, double, double)> sdf,
                             std::function<unsigned(double, double, double)> target_level,
                             double band, bool balance) {
    return amr::refineToSdfGraded(
        t_, geo_, [&](const Vec<3>& p) { return sdf(p[0], p[1], p[2]); },
        [&](const Vec<3>& p) { return target_level(p[0], p[1], p[2]); }, band, balance);
  }
 
  // The plan's gap-width floor (§7 criterion 1) driving refine_to_sdf_graded in one call:
  // target level = coarsest L with n*h_L <= gap(x), clamped to [0, coarsest_level].
  Index refine_to_gap_floor(std::function<double(double, double, double)> sdf,
                            std::function<double(double, double, double)> gap,
                            unsigned coarsest_level, double n, double band, bool balance) {
    auto tgt = amr::gapFloorTarget<3>([&](const Vec<3>& p) { return gap(p[0], p[1], p[2]); },
                                      geo_.hMax(), coarsest_level, n);
    return amr::refineToSdfGraded(
        t_, geo_, [&](const Vec<3>& p) { return sdf(p[0], p[1], p[2]); }, tgt, band, balance);
  }
 
  // Split a single leaf by index; returns True if it was split (level>0).
  bool refine_leaf(Index i) { return t_.refineLeaf(i); }
 
  // Enforce 2:1 (graded) balance; returns the number of refinements performed.
  Index balance() { return t_.balance2to1(); }
 
  // Internal (not bound to Python): give the Poisson solver the underlying octree + geometry.
  const BO& octreeRef() const { return t_; }
  const amr::AmrGeometry<3>& geoRef() const { return geo_; }
 
  // ---- solution-adaptive refinement ----
  // Löhner normalized-second-difference indicator E in [0,1] per leaf from a scalar field (N,);
  // large E marks steep features to refine, small E marks smooth regions to coarsen.
  nb::ndarray<nb::numpy, double> lohner_indicator(DArray field, double eps) {
    return vecToArray(amr::lohnerIndicator(t_, asField(t_, field, "lohner_indicator"), eps));
  }
 
  // One solution-adaptive step driven by the Löhner indicator: refine leaves where the indicator
  // exceeds refine_thresh (down to finest_level), coarsen sibling groups all below coarsen_thresh,
  // restore 2:1 balance, and conservatively remap `field` (N,) onto the new mesh. MUTATES the
  // octree in place (num_leaves/centers/... then reflect the new mesh) and returns the remapped
  // field (M,).
  nb::ndarray<nb::numpy, double> adapt(DArray field, double refine_thresh, double coarsen_thresh,
                                       unsigned finest_level, double eps, bool linear) {
    auto fv = asField(t_, field, "adapt");
    auto r = amr::adapt(t_, fv, refine_thresh, coarsen_thresh, finest_level, eps, linear);
    t_ = std::move(r.octree);
    return vecToArray(std::move(r.field));
  }
 
  // ---- output ----
  // Write the octree + a per-leaf scalar field (N,) as a VTK UnstructuredGrid (.vtu, ASCII).
  void write_vtu(const std::string& path, const std::string& name, DArray field) {
    amr::writeVtu(path, t_, geo_, name, asField(t_, field, "write_vtu"));
  }
 
 private:
  BO t_;
  amr::AmrGeometry<3> geo_;
};
 
// ---- geometric-multigrid Poisson solver --------------------------------------------------------
 
// Cell-centered finite-volume Poisson solver (Lu = rhs) on an Octree, via a geometric multigrid
// V-cycle (peclet::amr::AmrMultigrid). The operator L is the conservative two-point FV
// Laplacian (negative-definite, suite sign convention); on a graded octree it is consistent and
// second-order in the bulk (first-order at coarse/fine faces). `periodic=True` solves the singular
// periodic problem (the constant null space is removed each cycle); the manufactured RHS b =
// apply(u_exact) is exactly mean-zero, so the residual drives to round-off. The hierarchy snapshots
// the octree at construction. (Cut-cell openness is handled by the flow solver, which sets it
// consistently.)
class Poisson : public Releasable {
 public:
  Poisson(const Octree& oct, bool periodic) : n_(oct.octreeRef().numLeaves()) {
    mg_.build(oct.octreeRef(), oct.geoRef().h0);
    mg_.setPeriodic(periodic);
  }
  void release() noexcept override { mg_ = amr::AmrMultigrid<3>{}; }
 
  Index num_leaves() const { return n_; }
  std::size_t num_levels() const { return mg_.numLevels(); }
 
  // L applied to u (the FV Laplacian); use it to manufacture a RHS b = apply(u_exact). (N,)->(N,).
  nb::ndarray<nb::numpy, double> apply(DArray u) {
    std::vector<double> out;
    mg_.op(0).applyLaplacian(toVec(u, "apply"), out);
    return vec1<double>(std::move(out));
  }
 
  // Volume-weighted L2 residual norm sqrt(sum V*(rhs - L u)^2).
  double residual(DArray u, DArray rhs) {
    std::vector<double> res;
    return mg_.op(0).residual(toVec(u, "residual"), toVec(rhs, "residual"), res);
  }
 
  // Solve L u = rhs with up to `cycles` V-cycles (pre/post smoothing sweeps each), starting from
  // x0 (or 0). Stops early once the residual <= tol (tol<=0 disables). Returns
  // (u (N,), final_residual, cycles_done).
  nb::tuple solve(DArray rhs, std::optional<DArray> x0, int cycles, int pre, int post, double tol) {
    std::vector<double> rv = toVec(rhs, "solve");
    std::vector<double> u(static_cast<std::size_t>(n_), 0.0);
    if (x0)
      u = toVec(*x0, "solve");
    std::vector<double> res;
    double r = mg_.op(0).residual(u, rv, res);
    int done = 0;
    for (int k = 0; k < cycles; ++k) {
      mg_.vcycle(0, u, rv, pre, post);
      r = mg_.op(0).residual(u, rv, res);
      ++done;
      if (tol > 0.0 && r <= tol)
        break;
    }
    return nb::make_tuple(vec1<double>(std::move(u)), r, done);
  }
 
 private:
  std::vector<double> toVec(const DArray& a, const char* who) const {
    if (static_cast<Index>(a.shape(0)) != n_)
      throw std::runtime_error(std::string(who) + ": length != num_leaves");
    const double* p = a.data();
    return std::vector<double>(p, p + static_cast<std::size_t>(n_));
  }
 
  amr::AmrMultigrid<3> mg_;
  Index n_;
};
 
// ---- collocated incompressible flow ------------------------------------------------------------
 
// Collocated (cell-centered) incompressible Stokes / Navier-Stokes step on an Octree with a
// cut-cell immersed boundary (peclet::amr::AmrFlow). Each step() is: an implicit
// backward-Euler viscous momentum predictor with no-slip (u=0) Dirichlet cut-cell IBM on the SDF
// solid, then the Almgren-Bell-Colella approximate projection in incremental-rotational form
// (openness-weighted pressure Poisson). Stokes by default; set_advection(True) adds explicit
// high-order (SOU / Koren-TVD) momentum advection. Driven by a body force; iterate step() to steady
// state. Velocities/pressure are per-leaf (num_leaves,) in Z-order slots. The octree is borrowed by
// reference (kept alive for the Flow's lifetime).
//
// The classic cut band is uniformly finest (the ghost closures reach ±2 same-level cells);
// set_ghost_sampled(True) is what admits cut cells at several levels. Away from the wall the
// advection and the C/F schemes handle 2:1 faces themselves (docs/amr_cf_convective.md).
class DistributedOctree;  // the MPI wrapper, defined below (Flow's distributed constructor)
 
class Flow : public Releasable {
 public:
  Flow(const Octree& oct, double rho, double mu, double dt) : n_(oct.octreeRef().numLeaves()) {
    flow_.init(oct.octreeRef(), oct.geoRef().h0, oct.geoRef().origin);
    flow_.setDensity(rho);
    flow_.setViscosity(mu);
    flow_.setDt(dt);
    const auto& t = oct.octreeRef();
    for (int d = 0; d < 3; ++d) {
      extent_[d] = static_cast<double>(t.brick()[d]) * static_cast<double>(Index(1) << t.lmax()) *
                   oct.geoRef().h0[d];
      origin_[d] = oct.geoRef().origin[d];
    }
  }
  // Distributed (mpi4py): run this solver on one ORB block of a DistributedOctree — the whole
  // step then executes multi-rank through the ±2 LeafHalo (docs/amr_distributed_flow.md).
  // Defined after DistributedOctree below; keep_alive on the binding pins the octree.
  Flow(DistributedOctree& d, double rho, double mu, double dt);
  // Distributed load rebalance: weighted-ORB re-decomposition migrating u/p with the leaves +
  // full rebuild (collective; distributed constructor only).
  void rebalance_mpi(std::function<double(double, double, double)> sdf) {
    nb::gil_scoped_release nogil;  // see set_solid
    flow_.rebalanceMpi([&](const Vec<3>& p) { return sdf(p[0], p[1], p[2]); });
    n_ = flow_.numLeaves();
  }
  void release() noexcept override { flow_ = amr::AmrFlow<>{}; }
 
  Index num_leaves() const { return n_; }
  void set_body_force(double fx, double fy, double fz) { flow_.setBodyForce(fx, fy, fz); }
  void set_advection(bool on) { flow_.setAdvection(on); }
  void set_ghost_projection(bool on, int matrix_order, int rhs_order) {
    flow_.setGhostProjection(on, matrix_order, rhs_order);
  }
  void set_ghost_sampled(bool on, double rho, long max_samples) {
    flow_.setGhostSampled(on, rho, max_samples);
  }
  void set_cf_scheme(int scheme) { flow_.setCfScheme(scheme); }
 
  // Adaptivity during a run: snapshot -> (externally mutate the Octree: adapt / refine_to_* /
  // balance) -> rebuild + conservative field transfer. num_leaves is refreshed.
  void begin_adapt() { flow_.beginAdapt(); }
  void set_pressure(nb::ndarray<double, nb::c_contig> h) {
    flow_.setPressure(std::vector<double>(h.data(), h.data() + h.shape(0)));
  }
  void set_dt(double dt) { flow_.setDt(dt); }
  void finish_adapt(std::function<double(double, double, double)> sdf) {
    nb::gil_scoped_release nogil;  // see set_solid
    flow_.finishAdapt([&](const Vec<3>& p) { return sdf(p[0], p[1], p[2]); });
    n_ = flow_.numLeaves();
  }
  void set_advection_scheme(int s) { flow_.setAdvectionScheme(s); }
  void set_implicit_advection(bool on) { flow_.setImplicitAdvection(on); }
 
  // Build the cut-cell operators from a signed-distance callable f(x,y,z) (>0 in fluid, <0 in
  // solid), and zero the velocity / pressure fields. Call before stepping; re-call to change the
  // geometry.
  void set_solid(std::function<double(double, double, double)> sdf) {
    // The host operator builders sample the SDF from Kokkos host parallel regions. nanobind's
    // std::function wrapper takes the GIL for every call, so with the calling thread still
    // holding it while it waits at the parallel region's barrier, the worker threads block
    // forever (measured: a hang under the OpenMP host backend at OMP_NUM_THREADS > 1). Release
    // it here; the callback re-acquires it per sample, which serialises the sampling but
    // cannot deadlock. The same holds for finish_adapt / rebalance_mpi, which rebuild through
    // setSolid.
    nb::gil_scoped_release nogil;
    flow_.setSolid([&](const Vec<3>& p) { return sdf(p[0], p[1], p[2]); });
  }
 
  // Native union-of-spheres SDF — the porous-media geometry evaluated in C++ instead of through
  // a Python callback. set_solid samples the SDF tens of times per leaf (operator build, overlay
  // classification at virtual positions, the openness probe), so at bed scale the callback, not
  // the solve, dominates: measured >1h43m of pure numpy on an 11.35M-leaf 180-sphere bed before
  // the GPU was touched once. `centers` is (M,3) and `radii` is (M,) or a single scalar-valued
  // (1,) array; `periodic` applies the minimum-image convention over the octree's own extent.
  void set_solid_spheres(nb::ndarray<const double, nb::c_contig> centers,
                         nb::ndarray<const double, nb::c_contig> radii, bool periodic) {
    if (centers.ndim() != 2 || centers.shape(1) != 3)
      throw std::runtime_error("set_solid_spheres: centers must be (M, 3)");
    const std::size_t m = centers.shape(0);
    if (radii.ndim() != 1 || (radii.shape(0) != m && radii.shape(0) != 1))
      throw std::runtime_error("set_solid_spheres: radii must be (M,) or (1,)");
    std::vector<double> cx(m), cy(m), cz(m), rr(m);
    const double* cp = centers.data();
    const double* rp = radii.data();
    const bool oneR = radii.shape(0) == 1;
    for (std::size_t i = 0; i < m; ++i) {
      cx[i] = cp[3 * i];
      cy[i] = cp[3 * i + 1];
      cz[i] = cp[3 * i + 2];
      rr[i] = oneR ? rp[0] : rp[i];
    }
    // Layer 2-for-core (suite/docs/archive/AMR_GEOMETRY_SETUP_REQUIREMENTS.md): the per-query
    // brute-force scan over all M spheres — measured 1.18 us/eval at M=180, 94% of the whole
    // setSolid — is retired for the candidate-grid + equal-radius query. BIT-IDENTICAL to the old
    // callback (value, not just sign) by construction and by the geom_scene_query ctest; the
    // templated setSolid inlines the query, so there is no std::function indirection left either.
    geom::SphereBedQuery q(std::move(cx), std::move(cy), std::move(cz), std::move(rr),
                           Vec<3>{origin_[0], origin_[1], origin_[2]},
                           Vec<3>{extent_[0], extent_[1], extent_[2]}, periodic);
    flow_.setSolid(q);
  }
 
  void set_outer_iterations(int n, double tol) { flow_.setOuterIterations(n, tol); }
  void set_pressure_tolerance(double rtol) { flow_.setPressureTol(rtol); }
  // The pressure DRIVER of the aperture projection: MG-PCG (true, the default) or the bounded
  // stationary V-cycle (false) -- flow's `set_pressure_pcg` spelling and tier (ROADMAP A'2).
  void set_pressure_pcg(bool on) { flow_.setPressurePCG(on); }
  void set_momentum_tolerance(double rtol) { flow_.setMomentumTol(rtol); }
  // §6.6/§11.4: what solves the coarsest pressure level, and where the ladder hands over to it --
  // flow's spellings and tier (NAMING.md §1.8). The MODE takes effect at once (the bottom is
  // re-decided on the hierarchy already built) and is kept for later set_solid calls; the extent
  // only at the next set_solid, where the ladder is built.
  void set_pressure_bottom(const std::string& mode) { flow_.setPressureBottom(mode); }
  void set_pressure_bottom_extent(long cells) {
    flow_.setPressureBottomExtent(static_cast<Index>(cells));
  }
  long pressure_bottom_extent() const { return static_cast<long>(flow_.pressureBottomExtent()); }
 
  // Advance one collocated projection step (Stokes, or NS if advection is on).
  void step(int mom_iters, int pres_iters) { flow_.step(mom_iters, pres_iters); }
 
  // Per-leaf velocity component c (0=x,1=y,2=z) -> (num_leaves,) float64.
  nb::ndarray<nb::numpy, double> velocity(int c) const {
    if (c < 0 || c > 2)
      throw std::runtime_error("velocity: component must be 0..2");
    const auto& v = flow_.velocity(c);
    return vec1<double>(std::vector<double>(v.begin(), v.end()));
  }
 
  // All three velocity components -> (num_leaves, 3) float64. Packed on-device into one interleaved
  // buffer, so the device->host boundary is crossed once (G6) instead of three times.
  nb::ndarray<nb::numpy, double> velocities() const { return vec2<double>(flow_.velocities(), 3); }
 
  // Per-leaf fluid mask (False inside the solid) -> (num_leaves,) bool.
  nb::ndarray<nb::numpy, bool> is_fluid() const {
    auto* buf = new std::vector<std::uint8_t>(static_cast<std::size_t>(n_));
    for (Index i = 0; i < n_; ++i)
      (*buf)[static_cast<std::size_t>(i)] = flow_.isFluid(i) ? 1 : 0;
    nb::capsule owner(buf,
                      [](void* p) noexcept { delete static_cast<std::vector<std::uint8_t>*>(p); });
    return nb::ndarray<nb::numpy, bool>(reinterpret_cast<bool*>(buf->data()),
                                        {static_cast<std::size_t>(n_)}, owner, {1});
  }
 
  // Volume-weighted L2 norm of the residual cell divergence (a projection-quality diagnostic).
  double divergence_norm() { return flow_.divNormL2(); }
 
  // Per-leaf pressure (incremental-rotational p) -> (num_leaves,) float64.
  nb::ndarray<nb::numpy, double> pressure() const { return vec1<double>(flow_.pressure()); }
  // ABC divergence-free FACE velocity, one value per CSR (sub)face -> (num_faces,) float64.
  nb::ndarray<nb::numpy, double> face_field() const { return vec1<double>(flow_.faceField()); }
  // Write velocity component c (0=x,1=y,2=z) from a (num_leaves,) array — ICs / restart /
  // warm-start.
  void set_velocity(int c, nb::ndarray<double, nb::c_contig> h) {
    flow_.setVelocity(c, std::vector<double>(h.data(), h.data() + h.shape(0)));
  }
  // Pressure projection only (no momentum) — project an externally-set velocity to divergence-free.
  void project(int pres_iters) { flow_.project(pres_iters); }
  amr::AmrFlow<>& engine() { return flow_; }
 
 private:
  amr::AmrFlow<> flow_;  // the canonical device (Kokkos) AMR flow
  Index n_;
  std::array<double, 3> extent_{0.0, 0.0, 0.0};  // world box size (min-image for set_solid_spheres)
  std::array<double, 3> origin_{0.0, 0.0, 0.0};  // world box origin (candidate-grid coverage)
};
 
// ---- distributed (MPI) octree ------------------------------------------------------------------
 
// The MPI octree: an ORB block decomposition of a global root grid, one BlockOctree per rank.
// Build it collectively (every rank constructs the same global geometry; ORB hands each rank a
// block), refine the local octree toward a global surface, restore cross-block 2:1 balance,
// load-rebalance leaves+fields onto a weighted ORB, gather face-neighbour field values across the
// owner-based halo, and read this rank's local leaf geometry / fields as numpy. Uses
// MPI_COMM_WORLD.
class DistributedOctree : public Releasable {
 public:
  DistributedOctree(std::array<long, 3> global_root, unsigned lmax, std::array<double, 3> origin,
                    std::array<double, 3> spacing, std::array<bool, 3> periodic) {
    amr::AmrGeometry<3> g;
    g.origin = {origin[0], origin[1], origin[2]};
    g.h0 = {spacing[0], spacing[1], spacing[2]};
    d_.init(IVec<3>{global_root[0], global_root[1], global_root[2]}, lmax, g,
            {periodic[0], periodic[1], periodic[2]}, MPI_COMM_WORLD);
  }
  void release() noexcept override { d_ = DO{}; }
 
  // ---- introspection ----
  int rank() const { return d_.rank(); }
  int size() const { return d_.size(); }
  Index num_leaves() const { return d_.local().numLeaves(); }
  unsigned lmax() const { return d_.lmax(); }
  std::array<double, 3> spacing() const {
    const auto& hv = d_.h0();
    return {hv[0], hv[1], hv[2]};
  }
  std::array<long, 3> cells() const {
    const auto& c = d_.globalFineSize();
    return {(long)c[0], (long)c[1], (long)c[2]};
  }
  std::array<double, 3> extent() const {
    const auto c = cells();
    const auto& hv = d_.h0();
    return {(double)c[0] * hv[0], (double)c[1] * hv[1], (double)c[2] * hv[2]};
  }
  std::array<double, 3> origin() const {
    const auto& o = d_.globalGeometry().origin;
    return {o[0], o[1], o[2]};
  }
  std::array<long, 3> block_origin_root() const {
    const auto& b = d_.blockOriginRoot();
    return {b[0], b[1], b[2]};
  }
  std::array<long, 3> block_brick() const {
    const auto& b = d_.blockBrick();
    return {b[0], b[1], b[2]};
  }
  std::array<long, 3> global_root_size() const {
    const auto& b = d_.globalRootSize();
    return {b[0], b[1], b[2]};
  }
 
  // World leaf geometry of THIS rank's block (centres in global world coordinates).
  nb::ndarray<nb::numpy, double> centers() const {
    return leafCenters(d_.local(), d_.localGeometry());
  }
  nb::ndarray<nb::numpy, double> sizes(int axis) const {
    return leafSizes(d_.local(), d_.localGeometry(), axis);
  }
  nb::ndarray<nb::numpy, std::int32_t> levels() const { return leafLevels(d_.local()); }
  nb::ndarray<nb::numpy, std::uint64_t> codes() const { return leafCodes(d_.local()); }
 
  // ---- refinement ----
  // Refine the local octree toward a GLOBAL sphere surface, then (if balance) restore cross-block
  // 2:1 balance collectively. Returns the local refinement count. Collective when balance=True.
  Index refine_to_sphere(std::array<double, 3> center, double radius, unsigned target_level,
                         double band, bool balance) {
    geom::Sphere s{{center[0], center[1], center[2]}, radius};
    auto lgeo = d_.localGeometry();
    Index n = amr::refineToSdf(
        d_.local(), lgeo, [&](const Vec<3>& p) { return s.eval(p); }, target_level, band,
        /*balance=*/false);
    if (balance)
      d_.balance();
    return n;
  }
 
  // Refine the local octree toward an arbitrary GLOBAL surface given as a Python SDF callable
  // f(x,y,z) (suite sign: <0 inside solid), then (if balance) restore cross-block 2:1 balance
  // collectively. The distributed analogue of the serial Octree.refine_to_sdf. Collective when
  // balance.
  Index refine_to_sdf(std::function<double(double, double, double)> sdf, unsigned target_level,
                      double band, bool balance) {
    auto lgeo = d_.localGeometry();
    Index n = amr::refineToSdf(
        d_.local(), lgeo, [&](const Vec<3>& p) { return sdf(p[0], p[1], p[2]); }, target_level,
        band, /*balance=*/false);
    if (balance)
      d_.balance();
    return n;
  }
 
  // Cross-block 2:1 balance (collective). Returns the local number of refinements performed.
  Index balance() { return d_.balance(); }
 
  // ---- load balancing ----
  // Re-decompose by leaf COUNT (weighted ORB) so each rank holds a near-equal share, migrating
  // leaves AND their fields. `fields` is (N,K) float64 (K per-leaf columns); returns this rank's
  // (M,K) columns after migration. Pure redistribution; the partition is updated in place.
  // Collective.
  nb::ndarray<nb::numpy, double> rebalance(DArray fields) {
    const Index n = d_.local().numLeaves();
    if (static_cast<Index>(fields.shape(0)) != n)
      throw std::runtime_error("rebalance: fields.shape[0] != num_leaves");
    const int K = static_cast<int>(fields.shape(1));
    const double* in = fields.data();
    std::vector<std::vector<double>> cols(static_cast<std::size_t>(K),
                                          std::vector<double>(static_cast<std::size_t>(n)));
    for (Index i = 0; i < n; ++i)
      for (int c = 0; c < K; ++c)
        cols[static_cast<std::size_t>(c)][static_cast<std::size_t>(i)] = in[i * K + c];
 
    d_.rebalance(cols);  // mutates the octree in place and swaps in the new columns
 
    const Index m = d_.local().numLeaves();
    std::vector<double> out(static_cast<std::size_t>(m) * static_cast<std::size_t>(K));
    for (Index i = 0; i < m; ++i)
      for (int c = 0; c < K; ++c)
        out[i * K + c] = cols[static_cast<std::size_t>(c)][static_cast<std::size_t>(i)];
    return vec2<double>(std::move(out), static_cast<std::size_t>(K));
  }
 
  // ---- halo ----
  // For each local leaf and each of the 6 faces, the neighbouring leaf's field value, gathered
  // across the owner-based halo. Returns (N,6) float64 laid out as [+x,-x,+y,-y,+z,-z]; domain
  // boundaries (no neighbour) carry `sentinel`. Collective. `field` is this rank's per-leaf (N,).
  nb::ndarray<nb::numpy, double> face_neighbor_gather(DArray field, double sentinel) {
    std::vector<double> f = asField(d_.local(), field, "face_neighbor_gather");
    std::vector<double> g = d_.faceNeighborGather(f, sentinel);
    return vec2<double>(std::move(g), 6);  // 2*Dim faces
  }
 
  // ---- solution-adaptive refinement ----
  // Löhner indicator per local leaf, evaluated across the owner-based halo so cross-block
  // neighbours contribute exactly as in a whole-domain solve. `field` is this rank's (num_leaves,).
  // Collective.
  nb::ndarray<nb::numpy, double> lohner_indicator(DArray field, double eps) {
    auto fv = asField(d_.local(), field, "lohner_indicator");
    return vecToArray(amr::lohnerIndicatorDistributed(d_, fv, eps));
  }
 
  // One distributed solution-adaptive step: refine/coarsen each block's local octree from the
  // Löhner indicator, restore cross-block 2:1 balance, and conservatively remap `field`
  // (num_leaves,) onto the new local mesh. MUTATES the octree in place (keeping ORB ownership);
  // returns the remapped local field (M,). Bit-identical across rank counts. Collective.
  nb::ndarray<nb::numpy, double> adapt(DArray field, double refine_thresh, double coarsen_thresh,
                                       unsigned finest_level, double eps, bool linear) {
    auto fv = asField(d_.local(), field, "adapt");
    return vecToArray(
        amr::distributedAdapt(d_, fv, refine_thresh, coarsen_thresh, finest_level, eps, linear));
  }
 
  // ---- output ----
  // Write THIS rank's local octree + a per-leaf scalar (N,) as a .vtu (one file per rank).
  void write_vtu(const std::string& path, const std::string& name, DArray field) {
    amr::writeVtu(path, d_.local(), d_.localGeometry(), name,
                  asField(d_.local(), field, "write_vtu"));
  }
 
  DO& ref() { return d_; }
 
 private:
  DO d_;
};
 
// Distributed Flow constructor (out of line: DistributedOctree is declared after Flow).
inline Flow::Flow(DistributedOctree& d, double rho, double mu, double dt)
    : n_(d.ref().local().numLeaves()) {
  flow_.initMpi(d.ref());
  flow_.setDensity(rho);
  flow_.setViscosity(mu);
  flow_.setDt(dt);
}
 
// ---- the diagnostics tier of Flow ------------------------------------------------------------
// suite/docs/QUALITY_PLAN.md D2: the PUBLIC surface of Flow is what a user needs to set up, run and
// read out a simulation; everything a developer uses to inspect, ablate or profile lives here, as
// `flow.diagnostics.<name>` — a view holding a reference to the Flow (no copies of state; the
// binding pins the Flow for the view's lifetime).
class FlowDiagnostics {
 public:
  explicit FlowDiagnostics(Flow& f) : f_(f) {}
  // Solver iteration counts from the last step (convergence / perf instruments).
  int last_mom_iters() const { return f_.engine().lastMomIters(); }
  int last_pres_iters() const { return f_.engine().lastPresIters(); }
  int last_outer_iters() const { return f_.engine().lastOuterIters(); }
  // L2 norm of the divergence of the ABC divergence-free FACE field (≈ the pressure-solve residual,
  // far below divergence_norm — including across 2:1 interfaces).
  double divergence_norm_face() { return f_.engine().divNormFace(); }
  // The C/F census (docs/amr_cf_flux_gate.md §6.5): a count, so a property (NAMING.md §1.2/§1.3).
  peclet::core::Index num_cf_cut_faces() const { return f_.engine().numCfCutFaces(); }
  // The face CSR topology `Flow.face_field()` is indexed by -> a dict of host-copied arrays (the
  // CSR itself, the per-face geometry, and the seam-reconstruction tables); the keys are listed in
  // the docstring below. A 2:1 sub-face is a slot whose two incident leaves have different
  // `Octree.levels()`; its centroid is the FINER leaf's face centre. Read-only and host-copied on
  // every call — a diagnostic, not a step-loop read-out.
  nb::dict face_topology() const {
    const auto t = f_.engine().faceTopology();
    nb::dict d;
    d["start"] = vec1<std::int64_t>(std::vector<std::int64_t>(t.start.begin(), t.start.end()));
    d["nbr"] = vec1<std::int64_t>(std::vector<std::int64_t>(t.nbr.begin(), t.nbr.end()));
    d["axis"] = vec1<std::int32_t>(std::vector<std::int32_t>(t.axis.begin(), t.axis.end()));
    d["dir"] = vec1<std::int32_t>(std::vector<std::int32_t>(t.dir.begin(), t.dir.end()));
    d["upup_i"] = vec1<std::int64_t>(std::vector<std::int64_t>(t.upupI.begin(), t.upupI.end()));
    d["upup_j"] = vec1<std::int64_t>(std::vector<std::int64_t>(t.upupJ.begin(), t.upupJ.end()));
    d["raw_area"] = vec1<double>(t.rawArea);
    d["dist"] = vec1<double>(t.dist);
    d["alpha"] = vec1<double>(t.alpha);
    // Seam reconstruction (docs/amr_cf_convective.md §5.2). Empty arrays when no tables exist.
    auto i64 = [](const std::vector<peclet::core::Index>& v) {
      return vec1<std::int64_t>(std::vector<std::int64_t>(v.begin(), v.end()));
    };
    d["seam"] = i64(t.seam);
    d["samp_i"] = i64(t.sampI);
    d["samp_j"] = i64(t.sampJ);
    d["uu_rec_i"] = i64(t.uuRecI);
    d["uu_rec_j"] = i64(t.uuRecJ);
    d["d1_i"] = vec1<double>(t.d1I);
    d["d1_j"] = vec1<double>(t.d1J);
    d["rec_start"] = i64(t.recStart);
    d["rec_dist"] = vec1<double>(t.recDist);
    d["rec_cell"] = i64(t.recCell);
    d["rec_w"] = vec1<double>(t.recW);
    d["cell_center"] = vec2<double>(std::vector<double>(t.center), 3);
    return d;
  }
  // The pressure multigrid's ladder AS BUILT on this rank: per-level leaf counts, level 0 first
  // (docs/amr_mg_depth.md §6.7). Levels below the root brick are lifted levels. Read it against
  // `peclet.amr.predict_hierarchy`, never against a literal.
  std::vector<peclet::core::Index> pressure_mg_levels() const {
    return f_.engine().pressureMgLevels();
  }
  // What actually solves the coarsest pressure level: "jacobi" | "amg", with "+tail" when the
  // replicated stage moved it (docs/amr_mg_depth.md §6.6/§6.7).
  std::string pressure_mg_bottom() const { return f_.engine().pressureMgBottom(); }
  // The seam-reconstruction census (docs/amr_cf_convective.md §5.2), this rank's.
  peclet::core::Index num_seam_sample_records() const { return f_.engine().numSeamSampleRecords(); }
  peclet::core::Index num_seam_layer_records() const { return f_.engine().numSeamLayerRecords(); }
  // Solver-internals and ablation switches (the production path needs none of them).
  void set_momentum_mg(bool on) { f_.engine().setMomentumMG(on); }
  void set_momentum_gs(bool on) { f_.engine().setMomentumGS(on); }
  void set_velocity_mg_staircase(bool on) { f_.engine().setVelocityMGStaircase(on); }
  void set_momentum_mg_solver(bool on) { f_.engine().setMomentumMGSolver(on); }
  void set_ghost_gradient(bool on) { f_.engine().setGhostGradient(on); }
  void set_aperture_order(int order) { f_.engine().setApertureOrder(order); }
  void set_uf_advection(bool on) { f_.engine().setUfAdvection(on); }
  void set_seam_reconstruction(bool on) { f_.engine().setSeamReconstruction(on); }
 
 private:
  Flow& f_;
};
 
// ---- the members Octree and DistributedOctree share, bound ONCE (QUALITY_PLAN G.5) --------------
// Both wrappers expose the same leaf read-outs, refinement / balance / adapt entry points and VTU
// writer with identical signatures (the distributed one acts on THIS rank's block and is
// collective where the docstring says so); the 14+ verbatim `.def` lines they used to carry each
// live here.
template <class T>
void bindOctreeCommon(nb::class_<T>& c) {
  c.def_prop_ro("cells", &T::cells,
                "Finest-level cell counts per axis (root*2**lmax; the GLOBAL grid on a "
                "DistributedOctree).")
      .def_prop_ro("extent", &T::extent, "Box side lengths in world units (cells*spacing).")
      .def_prop_ro("spacing", &T::spacing,
                   "Finest cell size (dx, dy, dz), per axis. Equal on a cubic octree.")
      .def_prop_ro("num_leaves", &T::num_leaves,
                   "Number of leaves (Z-order slots; this rank's on a DistributedOctree).")
      .def_prop_ro("lmax", &T::lmax, "Root-cell level (max refinement depth).")
      .def_prop_ro("origin", &T::origin, "Lower corner in world coordinates.")
      .def("centers", &T::centers,
           "Leaf world centres, (num_leaves, 3) float64 (global coordinates).")
      .def("sizes", &T::sizes, nb::arg("axis") = 0,
           "Leaf world widths along `axis`: spacing[axis]*2**level, (num_leaves,) float64. A leaf "
           "is a BOX, so `axis` selects which of the three widths (0 by default, which is THE "
           "width on a cubic octree).")
      .def("levels", &T::levels, "Leaf refinement levels, (num_leaves,) int32 (0 = finest).")
      .def("codes", &T::codes, "Leaf block-local Morton origin codes, (num_leaves,) uint64.")
      .def("refine_to_sphere", &T::refine_to_sphere, nb::arg("center"), nb::arg("radius"),
           nb::arg("target_level") = 0u, nb::arg("band") = 1.0, nb::arg("balance") = true,
           "Refine leaves the sphere surface passes through (plus `band` cells) down to "
           "target_level; optionally restore 2:1 balance (cross-block, collective, on a "
           "DistributedOctree). Returns the (local) number of refinements performed.")
      .def(
          "refine_to_sdf", &T::refine_to_sdf, nb::arg("sdf"), nb::arg("target_level") = 0u,
          nb::arg("band") = 1.0, nb::arg("balance") = true,
          "Refine toward an arbitrary signed-distance field given as a callable f(x,y,z)->distance "
          "(suite sign: <0 inside solid), down to target_level — rings / packed beds / any "
          "non-sphere geometry. Collective when balance=True on a DistributedOctree. Returns "
          "refinements performed.")
      .def("balance", &T::balance,
           "Enforce 2:1 graded balance to a fixpoint (cross-block and collective on a "
           "DistributedOctree); returns (this rank's) refinements performed.")
      .def(
          "lohner_indicator", &T::lohner_indicator, nb::arg("field"), nb::arg("eps") = 0.01,
          "Löhner normalized-second-difference feature indicator E in [0,1] per leaf from a scalar "
          "field (num_leaves,); large E = steep feature (refine), small = smooth (coarsen). On a "
          "DistributedOctree it is evaluated across the owner-based halo (collective).")
      .def("adapt", &T::adapt, nb::arg("field"), nb::arg("refine_thresh"),
           nb::arg("coarsen_thresh"), nb::arg("finest_level") = 0u, nb::arg("eps") = 0.01,
           nb::arg("linear") = true,
           "Solution-adaptive step (Löhner-driven): refine where the indicator > refine_thresh (to "
           "finest_level), coarsen sibling groups all < coarsen_thresh, 2:1-balance, and "
           "conservatively remap `field`. MUTATES the octree in place; returns the remapped field "
           "(M,). `linear` uses minmod-limited prolongation (else piecewise-constant). On a "
           "DistributedOctree: per block, cross-block balance, ORB ownership kept, bit-identical "
           "across rank counts (collective).")
      .def("write_vtu", &T::write_vtu, nb::arg("path"), nb::arg("name"), nb::arg("field"),
           "Write the octree (this rank's block on a DistributedOctree — one file per rank, "
           "combine in ParaView) + a per-leaf scalar field (num_leaves,) as a VTK UnstructuredGrid "
           "(.vtu, ASCII, one cell per leaf).");
}
 
}  // namespace peclet::amr::pybind
 
NB_MODULE(_amr, m) {
  using namespace peclet::amr::pybind;
  // The Flow path runs Kokkos kernels — initialise the device runtime on import (the backend/arch
  // is fixed by the prefix the module was built against); the release-then-finalize atexit hook,
  // finalize() and execution_space are the suite-wide pattern (teardown registry above). Per-leaf
  // arrays are host-vector-backed (no Views).
  peclet::core::python::install(m);
  m.attr("build_toolchain") = PECLET_AMR_BUILD_TOOLCHAIN;  // the byte gate compares like with like
  m.attr("__doc__") =
      "peclet.amr — adaptive mesh refinement: per-block Octree (serial) and DistributedOctree "
      "(MPI ORB) for the mesh, the geometric-multigrid Poisson solver, and the device (Kokkos) "
      "Flow cut-cell Stokes/Navier-Stokes solver. Build a graded octree, refine to an SDF "
      "surface, read leaf geometry + per-leaf fields as numpy, load-rebalance, gather face "
      "neighbours, export VTU, and run the flow step on device.";
 
  m.def(
      "predict_hierarchy",
      [](std::array<long, 3> cells, unsigned lmax, int num_ranks, long bottom_extent) {
        const auto root = rootCellsOf(cells, lmax, "predict_hierarchy");
        if (num_ranks < 1)
          throw std::runtime_error("predict_hierarchy: num_ranks must be >= 1");
        if (bottom_extent < 1)
          throw std::runtime_error("predict_hierarchy: bottom_extent must be >= 1");
        const peclet::core::IVec<3> G{root[0], root[1], root[2]};
        const auto p = peclet::amr::predictPressureLadder<3>(G, lmax, num_ranks, bottom_extent);
        nb::list cellsL, kindL, extentL;
        for (const auto& lv : p.levels) {
          cellsL.append(lv.cells);
          kindL.append(peclet::amr::toString(lv.kind));
          extentL.append(nb::make_tuple(lv.extent[0], lv.extent[1], lv.extent[2]));
        }
        nb::dict d;
        d["cells"] = cellsL;
        d["extent"] = extentL;
        d["kind"] = kindL;
        d["num_levels"] = (long)p.levels.size();
        d["num_in_place"] = (long)p.numInPlace();
        d["tail"] = p.tail;
        nb::list stagesL;
        for (auto k : p.stages)
          stagesL.append(peclet::amr::toString(k));
        d["stages"] = stagesL;
        d["bottom"] = p.bottomName();
        return d;
      },
      nb::arg("cells"), nb::arg("lmax") = 0u, nb::arg("num_ranks") = 1,
      nb::arg("bottom_extent") = 4,
      "The pressure-multigrid ladder a Flow builds on `cells` FINEST cells per axis at tree depth "
      "`lmax` over `num_ranks` ranks (docs/amr_mg_depth.md §6.2, §6.7) — a pre-flight tool. A "
      "pure function: it builds no mesh and needs no MPI; it re-runs the ladder rule on the ORB "
      "a freshly constructed DistributedOctree would produce (not a partition after "
      "`DistributedOctree.rebalance` / `Flow.rebalance_mpi`). Below the root brick the ladder "
      "halves the grid while every rank's block stays even; where it stops above "
      "`bottom_extent` cells per axis on more than one rank the level MOVES (a stage): onto fewer "
      "ranks by a sibling merge or a repartition when the Flow's stage policy is on "
      "(docs/amr_mg_depth.md WO4b), where the ladder lifts again, else gathered onto every rank "
      "(the replicated tail) and continued there; the coarsest level is then solved "
      "by damped-Jacobi sweeps if it has at most `bottom_extent` cells per axis, else by the "
      "agglomerated GraphAMG-PCG bottom. The prediction is that of the default "
      "`Flow.set_pressure_bottom('auto')`.\n\n"
      "Returns a dict, levels finest first: `extent` the level's global grid in CELLS per axis "
      "(a count, not a length), `cells` its total cell count as if the mesh were uniform (so it "
      "equals `Flow.diagnostics.pressure_mg_levels` leaf for leaf only for a uniform mesh at "
      "num_ranks=1), `kind` ('octree' at or above the root brick, 'lifted' below it in place, "
      "'tail' / 'sibling' / 'repartition' on a stage's continued ladder -- a stage's first level "
      "is the SAME grid as the level above it, moved rather than coarsened, exactly as "
      "`pressure_mg_levels` lists it), `num_levels`, `num_in_place` (levels above the first "
      "stage), `stages` (the stage kinds, outermost first), `tail` (bool: a replicated tail "
      "engages) and `bottom` ('jacobi' or 'amg', suffixed '+tail' / '+sibling' / '+repartition' "
      "per stage, innermost first) — the spelling `Flow.diagnostics.pressure_mg_bottom` "
      "reports.\n\n"
      "`lmax` is the number of octree coarsenings THE MESH supports, i.e. the tree's lmax for a "
      "mesh refined to level 0 somewhere. An UNREFINED Octree(cells, lmax=k>0) is the same mesh "
      "as Octree(cells/2**k, lmax=0) — all its leaves are root cells — and must be predicted that "
      "way. In general pass depth = tree.lmax - levels().min() and cells = root * 2**depth. "
      "`bottom_extent` (default 4, the solver's default) is the "
      "`Flow.pressure_bottom_extent` the run will use; pass the same value.");
 
  m.def(
      "spacing_from_extent",
      [](std::array<double, 3> extent, std::array<long, 3> root_cells, unsigned lmax) {
        const long f = (long)1 << lmax;
        return spacingFromExtent(extent, {root_cells[0] * f, root_cells[1] * f, root_cells[2] * f});
      },
      nb::arg("extent"), nb::arg("root_cells"), nb::arg("lmax"),
      "The finest cell width h0 = extent / (root_cells * 2**lmax) of a PHYSICAL domain — the "
      "one place an AMR spacing is computed, so no caller writes one (see "
      "suite/docs/PHYSICAL_UNITS_PLAN.md). Returns ONE number, so it raises when the extent does "
      "not give cubic cells — use `spacings_from_extent` for a box mesh.");
 
  m.def(
      "spacings_from_extent",
      [](std::array<double, 3> extent, std::array<long, 3> root_cells, unsigned lmax) {
        const long f = (long)1 << lmax;
        return spacingsFromExtent(extent,
                                  {root_cells[0] * f, root_cells[1] * f, root_cells[2] * f});
      },
      nb::arg("extent"), nb::arg("root_cells"), nb::arg("lmax"),
      "The finest cell size (dx, dy, dz) = extent / (root_cells * 2**lmax), PER AXIS. The "
      "anisotropic form of `spacing_from_extent`: it accepts any positive extent, because the "
      "octree's cells are boxes (core/docs/amr_anisotropic.md).");
 
  nb::class_<Octree> octree(
      m, "Octree",
      "Serial single-block adaptive octree with a world placement (origin + finest "
      "spacing per axis). Leaves are addressed in Z-order slot order; every per-leaf "
      "array is indexed by that slot.");
  octree
      .def(
          "__init__",
          [](Octree* self, std::array<long, 3> cells, unsigned lmax, std::array<double, 3> origin,
             std::optional<SpacingArg> spacing, std::optional<std::array<double, 3>> extent) {
            const auto root = rootCellsOf(cells, lmax, "Octree");
            new (self) Octree(root, lmax, origin, resolveSpacing(spacing, extent, cells, "Octree"));
          },
          nb::arg("cells"), nb::kw_only(), nb::arg("lmax") = 0u,
          nb::arg("origin") = std::array<double, 3>{0, 0, 0}, nb::arg("spacing") = nb::none(),
          nb::arg("extent") = nb::none(),
          "Build a uniform octree resolving `cells` FINEST cells per axis, refinable `lmax` levels "
          "(the root brick is cells/2**lmax root cells, which must divide exactly; lmax=0 is a "
          "uniform grid). `origin` is the block's lower corner. Every argument after `cells` is "
          "keyword-only.\n\n"
          "State the domain PHYSICALLY with `extent` = the box side lengths: the cell size is "
          "then extent/cells PER AXIS and the caller never writes a spacing. Alternatively give "
          "`spacing` (one number for cubic cells, or (dx, dy, dz)); default 1. Passing both "
          "raises.\n\n"
          "The cells are BOXES (core/docs/amr_anisotropic.md): any positive extent is accepted, "
          "the octree refines by 2 on every axis, and every level inherits the root aspect "
          "ratio. Read the three numbers back from `.spacing`.")
      .def("is_balanced", &Octree::is_balanced,
           "True iff every face-adjacent leaf pair differs by at most one level (2:1).")
      .def("find", &Octree::find, nb::arg("x"),
           "Index of the leaf containing world point x=(x,y,z), or -1 if outside the block.")
      .def("refine_to_sdf_graded", &Octree::refine_to_sdf_graded, nb::arg("sdf"),
           nb::arg("target_level"), nb::arg("band") = 2.0, nb::arg("balance") = true,
           "GRADED surface refinement (the mixed-level cut band, "
           "docs/amr_mixed_level_cut_band_plan.md §7): `target_level` is a callable "
           "f(x,y,z)->level giving the COARSEST acceptable level at a world point (0 = finest), "
           "so cut cells end up at SEVERAL levels — fine in throats/contacts, coarse on smooth "
           "caps. The band margin is measured in cells of the level being created (not in the "
           "finest spacing as in refine_to_sdf). Requires Flow.set_ghost_sampled(True) — the "
           "classic overlay contracts a uniform finest band and raises on these level jumps. "
           "Returns refinements performed.")
      .def("refine_to_gap_floor", &Octree::refine_to_gap_floor, nb::arg("sdf"), nb::arg("gap"),
           nb::arg("coarsest_level"), nb::arg("n") = 4.0, nb::arg("band") = 2.0,
           nb::arg("balance") = true,
           "refine_to_sdf_graded driven by the plan's GAP-WIDTH FLOOR (§7 criterion 1): the "
           "target level at a point is the coarsest L with n*h_L <= gap(x), clamped to "
           "[0, coarsest_level]. `gap` is the local fluid-gap proxy f(x,y,z)->width — for a "
           "sphere packing the two-closest-surfaces sum d1+d2; a medial-axis or peclet.pnm "
           "throat-radius field substitutes verbatim. n=4 per the M1/M2 measurements. This is "
           "the AMReX multi-valued-cell rule inverted: coarsening never merges or disconnects "
           "fluid.")
      .def("refine_leaf", &Octree::refine_leaf, nb::arg("i"),
           "Split leaf `i` into its 8 children; returns True if it was split (level>0).");
  bindOctreeCommon(octree);
 
  nb::class_<Poisson>(
      m, "Poisson",
      "Cell-centered finite-volume Poisson solver (L u = rhs) on an Octree, by a "
      "geometric-multigrid V-cycle. L is the conservative two-point FV Laplacian (suite sign). "
      "The hierarchy snapshots the octree at construction; per-leaf arrays are (num_leaves,) "
      "float64 in Z-order slots.")
      .def(
          nb::init<const Octree&, bool>(), nb::arg("octree"), nb::arg("periodic") = true,
          "Build the multigrid hierarchy from `octree`. periodic=True solves the singular periodic "
          "problem (constant null space removed each cycle).")
      .def_prop_ro("num_leaves", &Poisson::num_leaves, "Leaves on the finest level.")
      .def_prop_ro(
          "num_levels", &Poisson::num_levels,
          "Number of multigrid levels: the octree's own coarsenings down to the root brick, then "
          "levels BELOW it (the root lifted, docs/amr_mg_depth.md §6.1-§6.2) for as long as some "
          "axis has more than 4 cells and every axis stays even with at least 2 cells after "
          "halving. So a uniform lmax=0 mesh has a real hierarchy (64^3: 5 levels).")
      .def("apply", &Poisson::apply, nb::arg("u"),
           "L applied to u (the FV Laplacian); use b = apply(u_exact) to manufacture a RHS. "
           "(num_leaves,) -> (num_leaves,).")
      .def("residual", &Poisson::residual, nb::arg("u"), nb::arg("rhs"),
           "Volume-weighted L2 residual norm sqrt(sum V*(rhs - L u)^2).")
      .def(
          "solve", &Poisson::solve, nb::arg("rhs"), nb::arg("x0") = std::nullopt,
          nb::arg("cycles") = 20, nb::arg("pre") = 2, nb::arg("post") = 2, nb::arg("tol") = 0.0,
          "Solve L u = rhs with up to `cycles` V-cycles (pre/post Gauss-Seidel sweeps), from x0 or "
          "0, stopping once residual <= tol (tol<=0 disables). Returns (u (num_leaves,), "
          "final_residual, cycles_done).");
 
  nb::class_<Flow>(
      m, "Flow",
      "Collocated incompressible Stokes/Navier-Stokes step on an Octree with a cut-cell immersed "
      "boundary (no-slip on an SDF solid). step() = implicit viscous momentum predictor + "
      "Almgren-Bell-Colella rotational projection, whose pressure Poisson equation is solved by "
      "multigrid-preconditioned CG (the hierarchy continues below the octree's root brick; "
      "`diagnostics.pressure_mg_levels`). The box is TRIPLY PERIODIC and the only driving is a "
      "uniform body force (set_body_force); a wall is an immersed solid. Any consistent unit "
      "system: density [mass/length^3], dynamic viscosity [mass/(length*time)], dt [time], "
      "lengths those of the Octree's extent. Iterate step() to a steady state or march in time; "
      "velocities and pressure are per-leaf (num_leaves,) arrays in Z-order slots.")
      .def(nb::init<const Octree&, double, double, double>(), nb::arg("octree"),
           nb::arg("density") = 1.0, nb::arg("viscosity") = 1.0, nb::arg("dt") = 1e6,
           nb::keep_alive<1,
                          2>(),  // keep the octree alive for the Flow's lifetime (borrowed by ref)
           "Create a flow on `octree` with the given density, dynamic viscosity and time step "
           "(defaults 1, 1 and 1e6). A large dt drives straight to the steady (Stokes) solution. "
           "The octree is borrowed by reference and must outlive the Flow's use of it.")
      .def(nb::init<DistributedOctree&, double, double, double>(), nb::arg("octree"),
           nb::arg("density") = 1.0, nb::arg("viscosity") = 1.0, nb::arg("dt") = 1e6,
           nb::keep_alive<1, 2>(),
           "DISTRIBUTED (mpi4py, collective): run this solver on one ORB block of a "
           "DistributedOctree — the whole step (momentum, pressure, overlays, adaptivity) then "
           "executes multi-rank through the ±2 leaf ghost registry. Per-leaf arrays are this "
           "rank's LOCAL leaves. setSolid/step/project/begin_adapt/finish_adapt/rebalance_mpi "
           "are collective. np=1 is bit-identical to the single-rank constructor.")
      .def("rebalance_mpi", &Flow::rebalance_mpi, nb::arg("sdf"),
           "DISTRIBUTED: weighted-ORB load rebalance — migrates the leaves WITH the state (u, p) "
           "to the new owners and rebuilds every solver structure (collective). num_leaves is "
           "refreshed.")
      .def_prop_ro("num_leaves", &Flow::num_leaves, "Number of leaves.")
      .def("set_solid", &Flow::set_solid, nb::arg("sdf"),
           "Build the cut-cell operators from a signed-distance callable f(x,y,z) (>0 fluid, <0 "
           "solid) and zero the fields. Call before stepping; re-call to change the geometry.")
      .def("set_solid_spheres", &Flow::set_solid_spheres, nb::arg("centers"), nb::arg("radii"),
           nb::arg("periodic") = true,
           "set_solid for a UNION OF SPHERES, evaluated natively instead of through a Python "
           "callback — the porous-media geometry. `centers` is (M,3), `radii` is (M,) or (1,); "
           "`periodic` uses the minimum-image convention over the octree's own box extent. "
           "Prefer this to set_solid at bed scale: set_solid samples the SDF tens of times per "
           "leaf (operator build, overlay classification at virtual positions, the openness "
           "probe), so a Python callback dominates everything else — measured >1h43m of pure "
           "numpy on an 11.35M-leaf 180-sphere bed before the GPU ran a single kernel.")
      .def("set_body_force", &Flow::set_body_force, nb::arg("fx"), nb::arg("fy"), nb::arg("fz"),
           "Set the body force per unit volume [force/length^3] driving the flow, e.g. a mean "
           "pressure gradient -dp/dx along x. It is the only driving the periodic box has. "
           "Default (0, 0, 0).")
      .def("set_advection", &Flow::set_advection, nb::arg("on"),
           "Enable explicit momentum advection (Navier-Stokes); off = Stokes.")
      .def("set_ghost_projection", &Flow::set_ghost_projection, nb::arg("on"),
           nb::arg("matrix_order") = 2, nb::arg("rhs_order") = 2,
           "DEFAULT since 2026-08-25 (AUTO: ghost, with an aperture fallback + stderr notice when "
           "the finest band is too thin): the fluid-only constraint scheme — family-free, "
           "unconditionally stable, protocol-independent (flow's attractor-campaign verdicts; == "
           "flow's set_collocated_scheme('ghost')). FULL directional ghost-cell projection (the "
           "AMR port): binary-openness pressure operator + wall-anchored closure overlay on the "
           "finest-band rows, MG-preconditioned BiCGStab, ghost-closed divergence constraint; "
           "implies set_ghost_gradient. (matrix_order, rhs_order) closure orders: (2, 2) default "
           "and the only pair cleared for production — the (1, 2) mixed form is march-UNSTABLE "
           "above ~2000 spheres (flow hardening Phase A), kept callable for parity records only. "
           "Raises if the finest band is too thin (a closure would cross a 2:1 boundary). Call "
           "before set_solid.")
      .def("set_ghost_sampled", &Flow::set_ghost_sampled, nb::arg("on"), nb::arg("rho") = 2.2,
           nb::arg("max_samples") = 0L,
           "MIXED-LEVEL CUT BAND (docs/amr_mixed_level_cut_band_plan.md): allow cut cells at "
           "MULTIPLE octree levels — the finest-band contract is dropped. Chain entries that "
           "cross a 2:1 boundary become degree-2 LS virtual samples at the uniform closure "
           "positions (identity weights at same level, so a uniform finest band is BIT-IDENTICAL "
           "to set_ghost_sampled(False)); face classification uses the level-aware canonical "
           "openness; the momentum xi-row seam correction and the wall-aware C/F tangential "
           "fallback ride along. Implies the ghost projection (engages when the resolved scheme "
           "is ghost — the AUTO default or an explicit set_ghost_projection(True)). Distributed "
           "since 2026-08-30 (the clouds are a deterministic probe set through the leaf halo). "
           "`rho` is the least-squares cloud radius factor (rho = factor * max(h, H); 2.2 = the "
           "shipped behaviour) and `max_samples` the nearest-N candidate cap (0 = uncapped) — the "
           "M2a cloud-economy knobs, inert at their defaults; do not change them in production "
           "without the M2a table. Call before set_solid.")
      .def("set_pressure", &Flow::set_pressure, nb::arg("values"),
           "Write the accumulated rotational pressure from a (num_leaves,) array — restart, or "
           "re-accumulation policies after finish_adapt (at steady-state dt the transferred p is "
           "the load-bearing state; zeroing it after a coarsening adapt lets it re-accumulate "
           "cleanly).")
      .def("set_dt", &Flow::set_dt, nb::arg("dt"),
           "Change the time step. dt is BAKED INTO the momentum operator at build time "
           "(idiag = rho/dt), so a set_dt must be followed by set_solid or the operator stays "
           "stale. set_solid reallocates and zeroes u and p, so a dt SWITCH mid-march is: read "
           "velocity()/pressure() -> set_dt -> set_solid -> set_velocity()/set_pressure(). That "
           "sequence is the dt-cycling protocol of the attractor-family batteries.")
      .def("begin_adapt", &Flow::begin_adapt,
           "Snapshot the octree topology + (u, p) ahead of an external mesh mutation (adapt / "
           "refine_to_sphere / refine_to_sdf / balance on the SAME Octree object). Pair with "
           "finish_adapt.")
      .def("finish_adapt", &Flow::finish_adapt, nb::arg("sdf"),
           "Rebuild the solver on the mutated octree and conservatively transfer the snapshotted "
           "u and p onto it (minmod-limited linear remap). `sdf` is re-sampled on the new leaves "
           "(pass the same geometry callable as set_solid); keep the cut band at the finest "
           "level on the new mesh (re-run refine_to_sdf on the geometry band after a "
           "solution-driven adapt). The advecting face field restarts from the cell average for "
           "one step.")
      .def("set_cf_scheme", &Flow::set_cf_scheme, nb::arg("scheme"),
           "Coarse/fine (2:1) interface scheme. 1 = Martin-Cartwright tangential quadratic, THE "
           "DEFAULT since 2026-09-21: second order, applied to the momentum diffusion, the "
           "divergence constraint and the pressure gradients (the pressure matrix/MG stays "
           "standard, which does not move the steady solution). 0 = the standard two-point flux, "
           "first-order at level boundaries, kept only to reproduce pre-2026-09-21 graded results. "
           "The default is INERT BY GEOMETRY on a uniform or finest-band mesh -- no C/F faces, no "
           "delta -- and matters where a 2:1 interface sees tangential variation, which on a "
           "self-similar graded ladder is the difference between order 1.60 and 0.41. Works with "
           "both the aperture and the ghost projection. The seam reconstruction of the advected "
           "value "
           "(diagnostics.set_seam_reconstruction) rides on the quadratic scheme: with 0 no seam "
           "tables are built. Call before set_solid.")
      .def("set_advection_scheme", &Flow::set_advection_scheme, nb::arg("scheme"),
           "High-order advection flux: 0 = second-order upwind (default), 1 = Koren TVD.")
      .def(
          "set_implicit_advection", &Flow::set_implicit_advection, nb::arg("on"),
          "Implicit first-order-upwind deferred-correction advection (default on): unconditionally "
          "stable. Off = fully explicit high-order advection.")
      .def("set_outer_iterations", &Flow::set_outer_iterations, nb::arg("n"), nb::arg("tol") = 1e-6,
           "Picard outer iterations over the lagged advection per step (default 1).")
      .def("set_pressure_tolerance", &Flow::set_pressure_tolerance, nb::arg("rtol"),
           "Relative tolerance of the PRESSURE solve, for both drivers -- the MG-PCG default and "
           "the ghost projection's BiCGStab (flow shares its set_pressure_pcg tolerance with its "
           "ghost BiCGStab the same way). Default 1e-10, the value that was hard-coded until "
           "2026-09-21, so leaving it alone reproduces every earlier result. The iteration CAP is "
           "step()'s `pres_iters`; this is the accuracy it works to, and loosening it is the "
           "cheapest cost knob in the step.")
      .def("set_pressure_pcg", &Flow::set_pressure_pcg, nb::arg("on"),
           "The pressure-solve DRIVER of the aperture projection. on=True (THE DEFAULT) is "
           "MG-preconditioned conjugate gradients: one V-cycle of the pressure multigrid "
           "preconditions CG, the fluid-mean component of the residual is removed every iteration "
           "(an aperture cut-cell right-hand side is incompatible by exactly that component), and "
           "the solve stops at set_pressure_tolerance or at step()'s / project()'s `pres_iters`. "
           "on=False is the bounded stationary V-cycle: exactly `pres_iters` V-cycles, with no "
           "tolerance test and no mean removal, so on cut-cell geometry its residual STALLS at "
           "the incompatible component instead of converging. Both solve the same pressure "
           "equation, so where both converge they give the same projected velocity to the solve "
           "tolerance; the choice moves cost and robustness, not the answer. Leave it on: "
           "on=False exists to A/B the driver on a compatible problem. It has NO effect under the "
           "ghost projection -- the default scheme -- whose nonsymmetric operator is always "
           "solved by MG-preconditioned BiCGStab; it acts with set_ghost_projection(False) or "
           "where the default fell back to the aperture projection. flow's "
           "`Solver.set_pressure_pcg` is the same selector, but flow has no stationary driver "
           "(on=False raises there) and takes the cap and tolerance as its max_iter / rtol "
           "arguments. Takes effect at the next step()/project(); on a DISTRIBUTED Flow pass "
           "the same value on every rank.")
      .def("set_momentum_tolerance", &Flow::set_momentum_tolerance, nb::arg("rtol"),
           "Relative tolerance of the per-component MOMENTUM solve (BiCGStab, MG-preconditioned "
           "by default). Default 1e-8. At the large dt used for steady drag the momentum operator "
           "degrades toward a bare elliptic Laplacian and the solve gets expensive; this bounds "
           "the over-solve. The cap is step()'s `mom_iters`. (flow spells the same concept "
           "set_velocity_residual_tolerance -- the momentum/velocity divergence between the two "
           "codes is a ../docs/NAMING.md item, not settled here.)")
      .def("set_pressure_bottom", &Flow::set_pressure_bottom, nb::arg("mode"),
           "What solves the coarsest level of the pressure multigrid (docs/amr_mg_depth.md "
           "§6.6). A V-cycle converges at a mesh-independent rate only if its coarsest level is "
           "effectively solved, and 60 damped-Jacobi sweeps solve a level only up to ~4 cells per "
           "axis (8e-9 at 4, 8e-3 at 8, 0.6 at 25). 'auto' (THE DEFAULT) engages the agglomerated "
           "GraphAMG-PCG bottom iff the coarsest level still has more than "
           "`pressure_bottom_extent` cells on some axis -- i.e. where the ladder ran out on an odd "
           "or badly factored grid -- and keeps the sweeps otherwise; 'smoother' always sweeps; "
           "'agglomerated' always solves exactly. The choice changes the preconditioner, not the "
           "converged pressure (to the solve tolerance), and at the default extent it does not "
           "change the iteration count either (§11.4(a): identical counts in all nine measured "
           "'smoother'/'agglomerated' pairs); leave it at 'auto' unless you are comparing bottom "
           "solves. The same three strings and default as flow's `Solver.set_pressure_bottom`. "
           "Takes effect at once on the hierarchy already built and is kept for later set_solid "
           "calls. On a DISTRIBUTED Flow call it on every rank: the exact bottom lives in the "
           "replicated stage's continued ladder, which the call may build. Raises on any other "
           "string.")
      .def("set_pressure_bottom_extent", &Flow::set_pressure_bottom_extent, nb::arg("cells"),
           "Where the pressure ladder stops coarsening below the root brick and hands over to the "
           "bottom solve, in CELLS PER AXIS of the coarsest level -- a count, not a length: the "
           "limit is the bottom smoother's, which solves a level only up to ~4 cells per axis "
           "whatever the physical domain (../docs/NAMING.md §1.8; docs/amr_mg_depth.md "
           "§6.2/§11.4). The ladder stops once no axis has more than `cells` cells, and "
           "set_pressure_bottom('auto') engages the exact bottom where it stopped above it. "
           "Default 4, the value at which the 60 bottom sweeps are exact; it is measured rather "
           "than preferred (§11.4: 8 moves the iteration count by ~1 on one case in three; "
           "tests/study/amr_pressure_depth.py --sweep), so there is rarely a reason to change it. "
           "Changes the preconditioner, not the converged pressure (to the solve tolerance). "
           "Takes effect at the NEXT set_solid, where the ladder is built. Raises if cells < 1.")
      .def_prop_ro("pressure_bottom_extent", &Flow::pressure_bottom_extent,
                   "The bottom extent, in cells per axis (a count, not a length), that "
                   "set_pressure_bottom_extent last stored (default 4) -- the one the next "
                   "set_solid builds the ladder with (docs/amr_mg_depth.md §6.8).")
      .def("step", &Flow::step, nb::arg("mom_iters") = 100, nb::arg("pres_iters") = 60,
           "Advance one collocated projection step of length dt on device. `mom_iters` caps the "
           "momentum solve (BiCGStab, MG-preconditioned) and `pres_iters` the pressure solve "
           "(MG-PCG, or BiCGStab under the ghost projection); each stops earlier at its "
           "set_momentum_tolerance / set_pressure_tolerance. Collective on a distributed Flow.")
      .def("velocity", &Flow::velocity, nb::arg("component"),
           "Per-leaf velocity component (0=x,1=y,2=z), (num_leaves,) float64.")
      .def("velocities", &Flow::velocities,
           "All three velocity components, (num_leaves, 3) float64.")
      .def("is_fluid", &Flow::is_fluid,
           "Per-leaf fluid mask (False in the solid), (num_leaves,) bool.")
      .def("divergence_norm", &Flow::divergence_norm,
           "Volume-weighted L2 norm of the residual cell divergence (projection-quality "
           "diagnostic).")
      .def("pressure", &Flow::pressure,
           "Per-leaf pressure (incremental-rotational p), (num_leaves,) float64.")
      .def("face_field", &Flow::face_field,
           "ABC divergence-free FACE velocity, one value per CSR (sub)face (conservative flux / "
           "streamline post-processing).")
      .def("set_velocity", &Flow::set_velocity, nb::arg("component"), nb::arg("values"),
           "Write velocity component c (0=x,1=y,2=z) from a (num_leaves,) array — initial "
           "conditions, restart, or warm-start. Call before step()/project().")
      .def("project", &Flow::project, nb::arg("pres_iters") = 60,
           "Pressure projection only (no momentum solve) — project an externally-set velocity "
           "field to divergence-free. Returns nothing; read the result via "
           "velocity()/velocities().")
      .def_prop_ro(
          "diagnostics", [](Flow& f) { return FlowDiagnostics(f); }, nb::keep_alive<0, 1>(),
          "The developer tier (FlowDiagnostics, a view onto this Flow): iteration counts of the "
          "last step, the face-field divergence and topology, the C/F and seam-reconstruction "
          "censuses, the pressure-multigrid ladder and its bottom solver, and the "
          "solver-internals / ablation switches. Nothing here is needed to set up, run or read "
          "out a simulation; everything here has a production default.");
 
  nb::class_<FlowDiagnostics>(
      m, "FlowDiagnostics",
      "Developer instruments and ablation switches of a Flow, reached as `flow.diagnostics` "
      "(suite/docs/QUALITY_PLAN.md D2: the public Flow surface is what a user needs to set up, "
      "run and read out a simulation; this is what a developer uses to inspect or ablate it).")
      .def("last_mom_iters", &FlowDiagnostics::last_mom_iters,
           "Total momentum BiCGStab iterations (summed over the 3 components) of the last step.")
      .def("last_pres_iters", &FlowDiagnostics::last_pres_iters,
           "Pressure PCG iterations of the last step.")
      .def("last_outer_iters", &FlowDiagnostics::last_outer_iters,
           "Picard outer iterations actually run in the last step (1 unless "
           "set_outer_iterations(>1)).")
      .def("divergence_norm_face", &FlowDiagnostics::divergence_norm_face,
           "L2 norm of the APERTURE-weighted divergence of the ABC face field. MEANINGFUL ON THE "
           "APERTURE PATH ONLY (set_ghost_projection(False)), where it is the pressure-solve "
           "residual, far below divergence_norm — including across 2:1 interfaces. Under the "
           "GHOST projection (the DEFAULT) the solved constraint is this divergence PLUS an "
           "overlay delta that is a functional of the CELL velocities, not of the face field, so "
           "this norm omits it and reads O(1) on a perfectly healthy solve (34 at N=32, 184 at "
           "N=64 on the Z&H sphere, with the velocities matching peclet.flow to 1e-6). There, "
           "divergence_norm() is the residual you want. Unnormalized either way: it grows with "
           "resolution and velocity magnitude, so read it as a trend, never as an absolute.")
      .def_prop_ro("num_cf_cut_faces", &FlowDiagnostics::num_cf_cut_faces,
                   "How many 2:1 C/F sub-face slots of THIS RANK carry the standard two-point "
                   "face value because the quadratic one is withheld: both incident cells must be "
                   "REGULAR fluid (fluid and not cut) for the quadratic C/F scheme "
                   "(set_cf_scheme(1), the default) to apply there, so this counts the sub-faces "
                   "where a level boundary meets the wall (docs/amr_cf_flux_gate.md §6.5). 0 "
                   "on every uniform or finest-band mesh; on a graded mesh it measures the size "
                   "of the set that runs at the standard scheme's local order. Each sub-face "
                   "contributes two slots (one per incident cell) and the sum over ranks equals "
                   "the single-rank count.")
      .def("face_topology", &FlowDiagnostics::face_topology,
           "The face CSR topology face_field() is indexed by, as a dict of arrays: 'start' "
           "(num_leaves+1 row offsets, int64), 'nbr' (neighbour leaf per (sub)face, int64), "
           "'axis' (0/1/2, int32), 'dir' (+1/-1 from the owning cell toward the neighbour, "
           "int32), 'raw_area' (world area the ADVECTIVE flux uses -- the FINE area at a 2:1 "
           "sub-face, so a coarse face's four sub-faces sum to the coarse area), 'dist' (world "
           "centre distance, 1.5*h_fine at a 2:1 sub-face), 'alpha' (face openness in [0, 1]) and "
           "'upup_i' / 'upup_j' (the second upwind probes the SOU/Koren reconstruction samples, "
           "-1 where none). A 2:1 sub-face is a slot whose two incident leaves have different "
           "Octree.levels(); its centroid is the FINER leaf's face centre. Host-copied on every "
           "call -- a diagnostic, not a step-loop read-out. Under MPI a neighbour index >= "
           "num_leaves is a ghost slot of this rank's registry, whose world centre (like every "
           "slot's) is row `slot` of 'cell_center', an (num_leaves + num_ghost_cells, 3) array. "
           "The SEAM RECONSTRUCTION tables of docs/amr_cf_convective.md come with it: 'seam' (one "
           "descriptor id per face slot, -1 = a plain slot that takes the ordinary SOU/Koren "
           "line), the per-descriptor 'samp_i' / 'samp_j' (the tangential-sample record when i / "
           "j is the COARSE cell of a 2:1 sub-face, else -1), 'uu_rec_i' / 'uu_rec_j' (the "
           "upstream probe's record when the second upwind cell of i / j crosses a level, else "
           "-1) and 'd1_i' / 'd1_j' (world half width along the face axis), and the record CSR "
           "'rec_start', 'rec_cell', 'rec_w', 'rec_dist' (the world probe distance, used only "
           "where a record is an UPSTREAM probe). All empty with set_cf_scheme(0 = standard), "
           "which builds no tables. Rank-local under MPI.")
      .def_prop_ro("pressure_mg_levels", &FlowDiagnostics::pressure_mg_levels,
                   "Leaf count of every pressure-multigrid level, finest first "
                   "(docs/amr_mg_depth.md §6.7), as a list of int; empty before set_solid, which "
                   "is where the hierarchy is built. Levels below the root brick are LIFTED levels "
                   "— the same octree with its root halved — so a uniform mesh has a real "
                   "hierarchy rather than a single level. Under MPI the in-place levels are THIS "
                   "RANK's counts, the levels of a replicated tail (appended last) are GLOBAL "
                   "counts, identical on every rank, and the levels of a sibling-merge or "
                   "repartition stage are this rank's counts on its target block (0 on a rank "
                   "that owns none). Check it against "
                   "`peclet.amr.predict_hierarchy` (level count, and leaf for leaf on a "
                   "uniform mesh at np=1), never against a literal.")
      .def_prop_ro("pressure_mg_bottom", &FlowDiagnostics::pressure_mg_bottom,
                   "What solves the coarsest pressure level of the hierarchy AS BUILT: 'jacobi' "
                   "(60 damped-Jacobi sweeps, effectively exact at <= 4 cells per axis: the "
                   "slowest mode falls by 8e-9) or 'amg' (the agglomerated GraphAMG-PCG solve), "
                   "with the suffix '+tail' when the coarsest level was gathered onto every rank "
                   "by the replicated stage, '+sibling' / '+repartition' when a level moved onto "
                   "fewer ranks (one suffix per stage, innermost first; docs/amr_mg_depth.md "
                   "§6.5-§6.7, WO4b). Which one runs "
                   "follows `Flow.set_pressure_bottom` (default 'auto': the exact bottom engages "
                   "only where the ladder ran out above `Flow.pressure_bottom_extent` cells per "
                   "axis). 'jacobi' before set_solid.")
      .def_prop_ro("num_seam_sample_records", &FlowDiagnostics::num_seam_sample_records,
                   "How many tangential-sample records the seam reconstruction built on THIS RANK "
                   "-- one per 2:1 sub-face pair that passes the C/F face gate (both cells regular "
                   "fluid). 0 on a uniform mesh and with the standard C/F scheme.")
      .def_prop_ro("num_seam_layer_records", &FlowDiagnostics::num_seam_layer_records,
                   "How many face-layer records the seam reconstruction built on THIS RANK -- one "
                   "per (coarse cell, face) whose far side is refined and whose four fine cells "
                   "are fluid (the case-3 upstream probe).")
      .def("set_momentum_mg", &FlowDiagnostics::set_momentum_mg, nb::arg("on"),
           "Use the Galerkin velocity multigrid as the momentum solve preconditioner (default on; "
           "makes the momentum solve scale with resolution). Call before set_solid.")
      .def("set_momentum_gs", &FlowDiagnostics::set_momentum_gs, nb::arg("on"),
           "Use the symmetric multicolour Gauss-Seidel smoother in the momentum multigrid (default "
           "off = weighted Jacobi). Call before set_solid.")
      .def("set_velocity_mg_staircase", &FlowDiagnostics::set_velocity_mg_staircase, nb::arg("on"),
           "Use the rediscretised staircase velocity-MG instead of Galerkin (default off).")
      .def("set_momentum_mg_solver", &FlowDiagnostics::set_momentum_mg_solver, nb::arg("on"),
           "Solve the momentum predictor with the velocity-MG as the solver (no Krylov), mirroring "
           "flow's velocity solve (default off = BiCgStab with the MG as preconditioner).")
      .def("set_ghost_gradient", &FlowDiagnostics::set_ghost_gradient, nb::arg("on"),
           "Directional ghost cell-gradient on cut cells for the pressure predictor and the "
           "projection's cell correction (2nd-order one-sided, never reads decoupled solid "
           "pressure — removes the gauge-dependent O(1/h) cut-cell gradient error of the plain "
           "ABC gradient). The ghost projection (the default) implies it; this switch matters for "
           "the aperture fallback only. Call before set_solid.")
      .def("set_aperture_order", &FlowDiagnostics::set_aperture_order, nb::arg("order"),
           "Aperture estimator for the (fallback) aperture projection: 2 = analytic "
           "marching-squares (DEFAULT since 2026-08-26), 1 = legacy one-sample model. "
           "Call before set_solid.")
      .def("set_uf_advection", &FlowDiagnostics::set_uf_advection, nb::arg("on"),
           "ABLATION. The advecting velocity of the momentum advection: the projected, "
           "divergence-free face field uf (ON, the shipped Almgren-Bell-Colella scheme) or, with "
           "on=False, the un-projected 1/2(u_i+u_j) cell->face average that the first step uses "
           "before any projection has run. It is the ONE discretization difference between this "
           "solver and peclet.flow's SolverColocated on a uniform grid; turning it off makes the "
           "two agree to solver tolerance (docs/amr_flow_uniform_parity.md).")
      .def("set_seam_reconstruction", &FlowDiagnostics::set_seam_reconstruction, nb::arg("on"),
           "ABLATION. Reconstruct the ADVECTED value with level-aware probes at every face whose "
           "upwind-side stencil crosses a 2:1 coarse/fine seam (default ON, "
           "docs/amr_cf_convective.md): the coarse upwind cell is tangentially sampled at the "
           "sub-face's own column, and an upstream probe at another octree level is taken at its "
           "TRUE distance (a coarser one as the same tangential sample, a finer one as the mean of "
           "the four face-layer children). Without it the sub-face value carries the coarse "
           "column's tangential offset -- an O(h) face error, so an O(1) local truncation on the "
           "fine side of the seam. on=False restores that, bit for bit, and may be flipped between "
           "steps; the tables are built in set_solid either way. Inert with advection off, with "
           "set_cf_scheme(0 = standard), and on any mesh with no 2:1 face. Leave it on in "
           "production: the switch exists for the A/B of docs/amr_cf_convective.md §12, where it "
           "cuts the one-step seam truncation 7.9e-3 -> 2.0e-3 at no measurable cost.");
 
  nb::class_<DistributedOctree> distributed(
      m, "DistributedOctree",
      "MPI octree: an ORB block decomposition of a global root grid (one BlockOctree per rank, "
      "over MPI_COMM_WORLD). Construct it collectively; "
      "refine/balance/rebalance/face_neighbor_gather are collective. Per-leaf arrays describe "
      "THIS rank's local block in global world coordinates.");
  distributed
      .def(
          "__init__",
          [](DistributedOctree* self, std::array<long, 3> cells, unsigned lmax,
             std::array<double, 3> origin, std::optional<SpacingArg> spacing,
             std::array<bool, 3> periodic, std::optional<std::array<double, 3>> extent) {
            const auto root = rootCellsOf(cells, lmax, "DistributedOctree");
            new (self) DistributedOctree(
                root, lmax, origin, resolveSpacing(spacing, extent, cells, "DistributedOctree"),
                periodic);
          },
          nb::arg("cells"), nb::kw_only(), nb::arg("lmax") = 0u,
          nb::arg("origin") = std::array<double, 3>{0, 0, 0}, nb::arg("spacing") = nb::none(),
          nb::arg("periodic") = std::array<bool, 3>{true, true, true},
          nb::arg("extent") = nb::none(),
          "Decompose a global grid of `cells` FINEST cells per axis (root cells = cells/2**lmax, "
          "each `lmax` levels deep) across the ranks of MPI_COMM_WORLD via ORB. `origin` places "
          "the global grid; `periodic` per axis. Every argument after `cells` is keyword-only.\n\n"
          "State the domain PHYSICALLY with `extent` = the global box side lengths; the cell size "
          "is then extent/cells PER AXIS. Alternatively give `spacing` (one number or (dx, dy, "
          "dz)); default 1. Passing both raises. Read the three numbers back from `.spacing`.")
      .def_prop_ro("rank", &DistributedOctree::rank, "This process's MPI rank.")
      .def_prop_ro("size", &DistributedOctree::size, "Number of ranks (blocks).")
      .def_prop_ro("block_origin_root", &DistributedOctree::block_origin_root,
                   "This rank's block lower corner, in global root-cell coordinates.")
      .def_prop_ro("block_brick", &DistributedOctree::block_brick,
                   "This rank's block size in root cells per axis.")
      .def_prop_ro("global_root_size", &DistributedOctree::global_root_size,
                   "Global grid size in root cells per axis.")
      .def("rebalance", &DistributedOctree::rebalance, nb::arg("fields"),
           "Re-decompose by leaf count (weighted ORB) and migrate leaves + their fields. `fields` "
           "is (num_leaves, K) float64; returns this rank's (M, K) columns after migration. Pure "
           "redistribution; the partition is updated in place (collective).")
      .def("face_neighbor_gather", &DistributedOctree::face_neighbor_gather, nb::arg("field"),
           nb::arg("sentinel") = 0.0,
           "For each local leaf, the field value across each of its 6 faces, gathered over the "
           "owner-based halo. `field` is (num_leaves,); returns (num_leaves, 6) laid out "
           "[+x,-x,+y,-y,+z,-z]; domain boundaries carry `sentinel` (collective).");
  bindOctreeCommon(distributed);
}