From 6526a3bd386d16bcc5a47b5c03263665271ca6a1 Mon Sep 17 00:00:00 2001 From: Omar Sobh Date: Thu, 17 Sep 2026 14:50:18 -0500 Subject: [PATCH] =?UTF-8?q?rtx-cfd=20embedded3=20item=201:=20the=20Poisson?= =?UTF-8?q?=20core=20as=20poisson/{problem,=20hierarchy,=20pcg}=20(?= =?UTF-8?q?=E2=89=A4=20385=20lines=20each;=20the=20periodic=20wrap=20built?= =?UTF-8?q?=20once=20into=20neighbour=20arrays);=20gate=201=20HELD:=20nz?= =?UTF-8?q?=3D1=20bit-identical=20to=20the=202D=20solver=20(lex=20+=20red-?= =?UTF-8?q?black,=20cached/uncached,=20iterations)=20and=20every=20extrusi?= =?UTF-8?q?on=20case=20bit-identical=20to=20the=20three=5Fd=20oracle,=20pl?= =?UTF-8?q?anes=20within=201.4e-11?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5.1 --- .../solvers/incompressible/embedded3/grid.rs | 83 ++++ .../solvers/incompressible/embedded3/mod.rs | 12 + .../embedded3/poisson/hierarchy.rs | 408 ++++++++++++++++++ .../incompressible/embedded3/poisson/mod.rs | 22 + .../incompressible/embedded3/poisson/pcg.rs | 294 +++++++++++++ .../embedded3/poisson/problem.rs | 256 +++++++++++ .../rtx-cfd/src/solvers/incompressible/mod.rs | 1 + .../tests/embedded3_poisson_identity.rs | 222 ++++++++++ 8 files changed, 1298 insertions(+) create mode 100644 crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/grid.rs create mode 100644 crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/mod.rs create mode 100644 crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/hierarchy.rs create mode 100644 crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/mod.rs create mode 100644 crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/pcg.rs create mode 100644 crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/problem.rs create mode 100644 crates/specialized/rtx-cfd/tests/embedded3_poisson_identity.rs diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/grid.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/grid.rs new file mode 100644 index 0000000..9f0af49 --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/grid.rs @@ -0,0 +1,83 @@ +//! The uniform Cartesian grid and its index conventions. + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct Grid { + pub nx: usize, + pub ny: usize, + pub nz: usize, + pub dx: f64, + pub dy: f64, + pub dz: f64, +} + +impl Grid { + /// Cubic cells of spacing `h`. + #[must_use] + pub fn cubic(nx: usize, ny: usize, nz: usize, h: f64) -> Self { + Self { + nx, + ny, + nz, + dx: h, + dy: h, + dz: h, + } + } + + #[inline] + #[must_use] + pub fn cells(&self) -> usize { + self.nx * self.ny * self.nz + } + + #[inline] + #[must_use] + pub fn cell(&self, k: usize, j: usize, i: usize) -> usize { + (k * self.ny + j) * self.nx + i + } + + /// `(k, j, i)` of a cell index. + #[inline] + #[must_use] + pub fn kji(&self, idx: usize) -> (usize, usize, usize) { + let nxy = self.nx * self.ny; + (idx / nxy, (idx % nxy) / self.nx, idx % self.nx) + } + + /// The u face west of cell `(k, j, i)`; `i = nx` is the east face of the last cell. + #[inline] + #[must_use] + pub fn uface(&self, k: usize, j: usize, i: usize) -> usize { + (k * self.ny + j) * (self.nx + 1) + i + } + + #[inline] + #[must_use] + pub fn vface(&self, k: usize, j: usize, i: usize) -> usize { + (k * (self.ny + 1) + j) * self.nx + i + } + + #[inline] + #[must_use] + pub fn wface(&self, k: usize, j: usize, i: usize) -> usize { + (k * self.ny + j) * self.nx + i + } + + #[inline] + #[must_use] + pub fn n_ufaces(&self) -> usize { + (self.nx + 1) * self.ny * self.nz + } + + #[inline] + #[must_use] + pub fn n_vfaces(&self) -> usize { + self.nx * (self.ny + 1) * self.nz + } + + #[inline] + #[must_use] + pub fn n_wfaces(&self) -> usize { + self.nx * self.ny * (self.nz + 1) + } +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/mod.rs new file mode 100644 index 0000000..34002e2 --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/mod.rs @@ -0,0 +1,12 @@ +//! `embedded3`: the 3D embedded solver, clean build (omni-cortex +//! `docs/embedded3_campaign.md`). A sharp-interface embedded wall on a +//! Cartesian grid, device-resident. The `three_d` module in history is the +//! oracle this one reproduces gate by gate; the 2D solver is never touched. +//! +//! Layout: cells `(k, j, i)` row-major, `cell = (k·ny + j)·nx + i`; +//! u faces on `(nx + 1)·ny·nz`, v on `nx·(ny + 1)·nz`, w on `nx·ny·(nz + 1)`. + +pub mod grid; +pub mod poisson; + +pub use grid::Grid; diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/hierarchy.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/hierarchy.rs new file mode 100644 index 0000000..d31bcbd --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/hierarchy.rs @@ -0,0 +1,408 @@ +//! The multigrid hierarchy: sanitised levels with explicit neighbour +//! arrays (the periodic wrap is data), aggregation by 2 per direction, the +//! symmetric Gauss–Seidel smoothers (lexicographic / red-black by +//! `(i + j + k) % 2`), the V-cycle with the ×2 coarse correction, and the +//! connected components of the active cells. + +use super::{COARSE_CORRECTION, COARSEST_SWEEPS, MAX_LEVELS, Problem}; +use crate::solvers::incompressible::poisson::{MgScalar, MgSmoother, MultigridParameters}; + +/// One level: the problem, its coefficients in the V-cycle scalar, the +/// active cells, the two colours, the neighbour arrays, the parent map. +#[derive(Clone)] +pub(crate) struct Level { + pub(crate) problem: Problem, + pub(crate) active: Vec, + pub(crate) ae: Vec, + pub(crate) aw: Vec, + pub(crate) an: Vec, + pub(crate) as_: Vec, + pub(crate) at: Vec, + pub(crate) ab: Vec, + pub(crate) ap: Vec, + pub(crate) cells: Vec, + pub(crate) red: Vec, + pub(crate) black: Vec, + /// Neighbour above / below per cell (`usize::MAX` = none). + pub(crate) top: Vec, + pub(crate) bot: Vec, + pub(crate) coarse_of: Vec, +} + +struct Work { + b: Vec, + x: Vec, + r: Vec, +} + +impl Level { + pub(crate) fn new(mut problem: Problem) -> Self { + let (nx, ny, nz) = (problem.nx, problem.ny, problem.nz); + let n = nx * ny * nz; + let ap: Vec = (0..n).map(|idx| problem.diagonal(idx)).collect(); + let active: Vec = (0..n) + .map(|idx| problem.active[idx] && ap[idx] > 0.0) + .collect(); + let mut top = vec![usize::MAX; n]; + let mut bot = vec![usize::MAX; n]; + for k in 0..nz { + for j in 0..ny { + for i in 0..nx { + let idx = problem.index(k, j, i); + let t = problem.top(idx, k); + let b = problem.bottom(idx, k); + if !active[idx] { + problem.ae[idx] = 0.0; + problem.aw[idx] = 0.0; + problem.an[idx] = 0.0; + problem.as_[idx] = 0.0; + problem.at[idx] = 0.0; + problem.ab[idx] = 0.0; + continue; + } + if !(i + 1 < nx && active[idx + 1]) { + problem.ae[idx] = 0.0; + } + if !(i > 0 && active[idx - 1]) { + problem.aw[idx] = 0.0; + } + if !(j + 1 < ny && active[idx + nx]) { + problem.an[idx] = 0.0; + } + if !(j > 0 && active[idx - nx]) { + problem.as_[idx] = 0.0; + } + match t { + Some(t) if active[t] => top[idx] = t, + _ => problem.at[idx] = 0.0, + } + match b { + Some(b) if active[b] => bot[idx] = b, + _ => problem.ab[idx] = 0.0, + } + } + } + } + let cells: Vec = (0..n).filter(|&idx| active[idx]).collect(); + let nxy = nx * ny; + let parity = |idx: usize| (idx % nx + (idx % nxy) / nx + idx / nxy) % 2; + let red: Vec = cells + .iter() + .copied() + .filter(|&idx| parity(idx) == 0) + .collect(); + let black: Vec = cells + .iter() + .copied() + .filter(|&idx| parity(idx) == 1) + .collect(); + let cast = |v: &[f64]| v.iter().map(|&x| T::from_f64(x)).collect::>(); + Self { + ae: cast(&problem.ae), + aw: cast(&problem.aw), + an: cast(&problem.an), + as_: cast(&problem.as_), + at: cast(&problem.at), + ab: cast(&problem.ab), + ap: cast(&ap), + problem, + active, + cells, + red, + black, + top, + bot, + coarse_of: Vec::new(), + } + } + + /// `Σ a_nb x_nb`: the 2D order (e, w, n, s) with t, b appended. + #[inline] + fn neighbour_sum(&self, x: &[T], idx: usize) -> T { + let nx = self.problem.nx; + let mut s = T::ZERO; + let ae = self.ae[idx]; + if ae != T::ZERO { + s += ae * x[idx + 1]; + } + let aw = self.aw[idx]; + if aw != T::ZERO { + s += aw * x[idx - 1]; + } + let an = self.an[idx]; + if an != T::ZERO { + s += an * x[idx + nx]; + } + let as_ = self.as_[idx]; + if as_ != T::ZERO { + s += as_ * x[idx - nx]; + } + let at = self.at[idx]; + if at != T::ZERO { + s += at * x[self.top[idx]]; + } + let ab = self.ab[idx]; + if ab != T::ZERO { + s += ab * x[self.bot[idx]]; + } + s + } + + pub(crate) fn apply(&self, x: &[T], y: &mut [T]) { + for &idx in &self.cells { + y[idx] = self.ap[idx] * x[idx] - self.neighbour_sum(x, idx); + } + } + + pub(crate) fn residual(&self, b: &[T], x: &[T], r: &mut [T]) -> T { + let mut l1 = T::ZERO; + for &idx in &self.cells { + let v = b[idx] - (self.ap[idx] * x[idx] - self.neighbour_sum(x, idx)); + r[idx] = v; + l1 += v.abs(); + } + l1 + } + + fn symmetric_gs(&self, b: &[T], x: &mut [T]) { + for &idx in &self.cells { + x[idx] = (b[idx] + self.neighbour_sum(x, idx)) / self.ap[idx]; + } + for &idx in self.cells.iter().rev() { + x[idx] = (b[idx] + self.neighbour_sum(x, idx)) / self.ap[idx]; + } + } + + fn symmetric_gs_rb(&self, b: &[T], x: &mut [T]) { + for list in [&self.red, &self.black, &self.black, &self.red] { + for &idx in list { + x[idx] = (b[idx] + self.neighbour_sum(x, idx)) / self.ap[idx]; + } + } + } + + fn smooth(&self, b: &[T], x: &mut [T], smoother: MgSmoother) { + match smoother { + MgSmoother::Lexicographic => self.symmetric_gs(b, x), + MgSmoother::RedBlack => self.symmetric_gs_rb(b, x), + } + } + + /// Galerkin coarsening by 2 in every direction (the 2D rule with k). + fn coarsen(&self) -> (Problem, Vec) { + let p = &self.problem; + let (nx, ny, nz) = (p.nx, p.ny, p.nz); + let nxc = (nx / 2).max(1); + let nyc = (ny / 2).max(1); + let nzc = (nz / 2).max(1); + let cx = |i: usize| (i / 2).min(nxc - 1); + let cy = |j: usize| (j / 2).min(nyc - 1); + let cz = |k: usize| (k / 2).min(nzc - 1); + let mut coarse = Problem::new(nxc, nyc, nzc); + coarse.periodic_z = p.periodic_z; + coarse.active.fill(false); + let mut coarse_of = vec![usize::MAX; nx * ny * nz]; + for &idx in &self.cells { + let (k, j, i) = (idx / (nx * ny), (idx % (nx * ny)) / nx, idx % nx); + let (ic, jc, kc) = (cx(i), cy(j), cz(k)); + let c = coarse.index(kc, jc, ic); + coarse_of[idx] = c; + coarse.active[c] = true; + coarse.extra_diag[c] += p.extra_diag[idx]; + if p.ae[idx] != 0.0 && cx(i + 1) != ic { + coarse.ae[c] += p.ae[idx]; + } + if p.aw[idx] != 0.0 && cx(i - 1) != ic { + coarse.aw[c] += p.aw[idx]; + } + if p.an[idx] != 0.0 && cy(j + 1) != jc { + coarse.an[c] += p.an[idx]; + } + if p.as_[idx] != 0.0 && cy(j - 1) != jc { + coarse.as_[c] += p.as_[idx]; + } + if p.at[idx] != 0.0 && cz(self.top[idx] / (nx * ny)) != kc { + coarse.at[c] += p.at[idx]; + } + if p.ab[idx] != 0.0 && cz(self.bot[idx] / (nx * ny)) != kc { + coarse.ab[c] += p.ab[idx]; + } + } + (coarse, coarse_of) + } +} + +/// The hierarchy: level 0 is the fine problem. +pub struct Hierarchy { + pub(crate) levels: Vec>, + work: Vec>, + sweeps: usize, + smoother: MgSmoother, +} + +impl Hierarchy { + pub fn build(problem: &Problem, params: &MultigridParameters) -> Self { + let mut levels = vec![Level::::new(problem.clone())]; + while levels.len() < MAX_LEVELS { + let fine = levels.last().expect("at least one level"); + if fine.cells.len() <= params.coarsest_cells.max(1) { + break; + } + let (coarse, coarse_of) = fine.coarsen(); + let coarse = Level::::new(coarse); + if coarse.cells.len() >= fine.cells.len() { + break; + } + let last = levels.len() - 1; + levels[last].coarse_of = coarse_of; + levels.push(coarse); + } + let work = levels + .iter() + .map(|l| { + let n = l.problem.nx * l.problem.ny * l.problem.nz; + Work { + b: vec![T::ZERO; n], + x: vec![T::ZERO; n], + r: vec![T::ZERO; n], + } + }) + .collect(); + Self { + levels, + work, + sweeps: params.smoother_sweeps.max(1), + smoother: params.smoother, + } + } + + pub fn depth(&self) -> usize { + self.levels.len() + } + + /// `z = M⁻¹ r`: one V-cycle from zero (the 2D sequence). + pub fn apply_preconditioner(&mut self, r: &[f64], z: &mut [f64]) { + let depth = self.levels.len(); + let levels = &self.levels; + for &idx in &levels[0].cells { + self.work[0].b[idx] = T::from_f64(r[idx]); + } + let smoother = self.smoother; + for l in 0..depth - 1 { + let (fine, coarse) = (&levels[l], &levels[l + 1]); + let (head, tail) = self.work.split_at_mut(l + 1); + let (wf, wc) = (&mut head[l], &mut tail[0]); + for &idx in &fine.cells { + wf.x[idx] = T::ZERO; + } + let Work { b, x, r } = wf; + for _ in 0..self.sweeps { + fine.smooth(b, x, smoother); + } + fine.residual(b, x, r); + for &idx in &coarse.cells { + wc.b[idx] = T::ZERO; + } + for &idx in &fine.cells { + let c = fine.coarse_of[idx]; + if coarse.active[c] { + wc.b[c] += r[idx]; + } + } + } + { + let bottom = &levels[depth - 1]; + let wb = &mut self.work[depth - 1]; + for &idx in &bottom.cells { + wb.x[idx] = T::ZERO; + } + let Work { b, x, .. } = wb; + for _ in 0..COARSEST_SWEEPS { + bottom.smooth(b, x, smoother); + } + } + for l in (0..depth - 1).rev() { + let fine = &levels[l]; + let (head, tail) = self.work.split_at_mut(l + 1); + let (wf, wc) = (&mut head[l], &tail[0]); + for &idx in &fine.cells { + wf.x[idx] += T::from_f64(COARSE_CORRECTION) * wc.x[fine.coarse_of[idx]]; + } + let Work { b, x, .. } = wf; + for _ in 0..self.sweeps { + fine.smooth(b, x, smoother); + } + } + for &idx in &levels[0].cells { + z[idx] = self.work[0].x[idx].to_f64(); + } + } +} + +/// Connected components of the active cells through non-zero faces, and +/// whether each is singular (no Dirichlet contribution). +#[derive(Clone)] +pub(crate) struct Components { + pub(crate) id: Vec, + pub(crate) members: Vec>, + pub(crate) singular: Vec, +} + +impl Components { + pub(crate) fn find(problem: &Problem, cells: &[usize]) -> Self { + let (nx, ny, nz) = (problem.nx, problem.ny, problem.nz); + let n = nx * ny * nz; + let mut id = vec![usize::MAX; n]; + let mut members = Vec::new(); + let mut singular = Vec::new(); + let mut stack = Vec::new(); + for &seed in cells { + if id[seed] != usize::MAX { + continue; + } + let c = members.len(); + let mut list = Vec::new(); + let mut has_dirichlet = false; + id[seed] = c; + stack.push(seed); + while let Some(idx) = stack.pop() { + list.push(idx); + if problem.extra_diag[idx] > 0.0 { + has_dirichlet = true; + } + let (k, j, i) = (idx / (nx * ny), (idx % (nx * ny)) / nx, idx % nx); + let mut visit = |nb: usize, coefficient: f64| { + if coefficient > 0.0 && problem.active[nb] && id[nb] == usize::MAX { + id[nb] = c; + stack.push(nb); + } + }; + if i + 1 < nx { + visit(idx + 1, problem.ae[idx]); + } + if i > 0 { + visit(idx - 1, problem.aw[idx]); + } + if j + 1 < ny { + visit(idx + nx, problem.an[idx]); + } + if j > 0 { + visit(idx - nx, problem.as_[idx]); + } + if let Some(t) = problem.top(idx, k) { + visit(t, problem.at[idx]); + } + if let Some(b) = problem.bottom(idx, k) { + visit(b, problem.ab[idx]); + } + } + members.push(list); + singular.push(!has_dirichlet); + } + Self { + id, + members, + singular, + } + } +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/mod.rs new file mode 100644 index 0000000..8e1da7f --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/mod.rs @@ -0,0 +1,22 @@ +//! The pressure Poisson solver: a seven-point masked variable-coefficient +//! operator, geometric multigrid by aggregation as the preconditioner of an +//! f64 conjugate gradient with a true-residual stop — the 2D `poisson.rs` +//! rule for rule. At `nz = 1` with zero z-coefficients the arithmetic is the +//! 2D solver's in the same order (gate 1: bit identity). + +mod hierarchy; +mod pcg; +mod problem; + +pub use hierarchy::Hierarchy; +pub use pcg::{PcgCache, solve_pcg, solve_pcg_cached}; +pub use problem::Problem; + +pub(crate) use hierarchy::{Components, Level}; +pub(crate) use pcg::OperatorKey; + +/// Symmetric GS sweeps on the coarsest level (the 2D value). +pub(crate) const COARSEST_SWEEPS: usize = 50; +/// The 2D `COARSE_CORRECTION`, proved dimension-independent there. +pub(crate) const COARSE_CORRECTION: f64 = 2.0; +pub(crate) const MAX_LEVELS: usize = 64; diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/pcg.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/pcg.rs new file mode 100644 index 0000000..204a499 --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/pcg.rs @@ -0,0 +1,294 @@ +//! The f64 conjugate gradient on a prepared operator (the 2D `run_pcg`, +//! line for line): absolute L1 true-residual stop with resynchronisation, +//! breakdown guard, per-component mean projection on singular components, +//! anchor shift on exit; the operator cache keyed on bit patterns. + +use super::{Components, Hierarchy, Level, Problem}; +use crate::solvers::incompressible::poisson::{ + MgPrecision, MgScalar, MgSmoother, MultigridParameters, PoissonSolution, +}; + +/// The operator part of a problem plus the hierarchy parameters (bit +/// patterns), to decide reuse of a prepared solver. +pub(crate) struct OperatorKey { + nx: usize, + ny: usize, + nz: usize, + periodic_z: bool, + active: Vec, + coefficients: Vec, + smoother_sweeps: usize, + coarsest_cells: usize, + smoother: MgSmoother, +} + +impl OperatorKey { + fn bits(problem: &Problem) -> impl Iterator + '_ { + problem + .ae + .iter() + .chain(&problem.aw) + .chain(&problem.an) + .chain(&problem.as_) + .chain(&problem.at) + .chain(&problem.ab) + .chain(&problem.extra_diag) + .map(|v| v.to_bits()) + } + + pub(crate) fn of(problem: &Problem, params: &MultigridParameters) -> Self { + Self { + nx: problem.nx, + ny: problem.ny, + nz: problem.nz, + periodic_z: problem.periodic_z, + active: problem.active.clone(), + coefficients: Self::bits(problem).collect(), + smoother_sweeps: params.smoother_sweeps, + coarsest_cells: params.coarsest_cells, + smoother: params.smoother, + } + } + + pub(crate) fn matches(&self, problem: &Problem, params: &MultigridParameters) -> bool { + self.nx == problem.nx + && self.ny == problem.ny + && self.nz == problem.nz + && self.periodic_z == problem.periodic_z + && self.smoother_sweeps == params.smoother_sweeps + && self.coarsest_cells == params.coarsest_cells + && self.smoother == params.smoother + && self.active == problem.active + && self.coefficients.iter().copied().eq(Self::bits(problem)) + } +} + +/// Everything the CG derives from the operator. +struct Prepared { + key: OperatorKey, + hier: Hierarchy, + fine: Level, + cells: Vec, + components: Components, +} + +impl Prepared { + fn build(problem: &Problem, params: &MultigridParameters) -> Self { + let hier = Hierarchy::::build(problem, params); + let fine = Level::::new(problem.clone()); + let cells = fine.cells.clone(); + let components = Components::find(problem, &cells); + Self { + key: OperatorKey::of(problem, params), + hier, + fine, + cells, + components, + } + } +} + +/// Prepared operators, reused while the operator is bit-identical. +#[derive(Default)] +pub struct PcgCache { + f64: Option>, + f32: Option>, +} + +/// CG preconditioned by one V-cycle (the 2D contract). +pub fn solve_pcg( + problem: &Problem, + p: &mut [f64], + params: &MultigridParameters, + tolerance: f64, + anchor: Option, +) -> PoissonSolution { + let t_entry = std::time::Instant::now(); + match params.precision { + MgPrecision::F64 => { + let mut prep = Prepared::::build(problem, params); + let setup_ns = t_entry.elapsed().as_nanos() as u64; + run_pcg(&mut prep, problem, p, params, tolerance, anchor, setup_ns) + } + MgPrecision::F32 => { + let mut prep = Prepared::::build(problem, params); + let setup_ns = t_entry.elapsed().as_nanos() as u64; + run_pcg(&mut prep, problem, p, params, tolerance, anchor, setup_ns) + } + } +} + +/// [`solve_pcg`] with the operator taken from `cache` on a hit. +pub fn solve_pcg_cached( + problem: &Problem, + p: &mut [f64], + params: &MultigridParameters, + tolerance: f64, + anchor: Option, + cache: &mut PcgCache, +) -> PoissonSolution { + match params.precision { + MgPrecision::F64 => { + solve_cached_with::(problem, p, params, tolerance, anchor, &mut cache.f64) + } + MgPrecision::F32 => { + solve_cached_with::(problem, p, params, tolerance, anchor, &mut cache.f32) + } + } +} + +fn solve_cached_with( + problem: &Problem, + p: &mut [f64], + params: &MultigridParameters, + tolerance: f64, + anchor: Option, + slot: &mut Option>, +) -> PoissonSolution { + let t_entry = std::time::Instant::now(); + let hit = slot + .as_ref() + .is_some_and(|prep| prep.key.matches(problem, params)); + if !hit { + *slot = Some(Prepared::::build(problem, params)); + } + let setup_ns = t_entry.elapsed().as_nanos() as u64; + let prep = slot.as_mut().expect("prepared"); + run_pcg(prep, problem, p, params, tolerance, anchor, setup_ns) +} + +fn run_pcg( + prep: &mut Prepared, + problem: &Problem, + p: &mut [f64], + params: &MultigridParameters, + tolerance: f64, + anchor: Option, + setup_ns: u64, +) -> PoissonSolution { + let n = problem.nx * problem.ny * problem.nz; + assert_eq!(p.len(), n, "p must have nx*ny*nz entries"); + debug_assert!( + problem.validate().is_ok(), + "invalid Problem: {:?}", + problem.validate() + ); + let Prepared { + hier, + fine, + cells, + components, + .. + } = prep; + let fine: &Level = fine; + let cells: &[usize] = cells; + let components: &Components = components; + let mut precond = |r: &[f64], z: &mut [f64]| hier.apply_preconditioner(r, z); + if cells.is_empty() { + return PoissonSolution { + iterations: 0, + residual: 0.0, + converged: true, + setup_ns, + iterate_ns: 0, + }; + } + let singular = components.singular.iter().any(|&s| s); + let project_mean = |v: &mut [f64]| { + for (c, members) in components.members.iter().enumerate() { + if !components.singular[c] { + continue; + } + let mean = members.iter().map(|&idx| v[idx]).sum::() / members.len() as f64; + for &idx in members { + v[idx] -= mean; + } + } + }; + let dot = |a: &[f64], b: &[f64]| cells.iter().map(|&idx| a[idx] * b[idx]).sum::(); + let l1 = |a: &[f64]| cells.iter().map(|&idx| a[idx].abs()).sum::(); + + let mut b = vec![0.0; n]; + for &idx in cells { + b[idx] = problem.rhs[idx]; + } + project_mean(&mut b); + + let mut r = vec![0.0; n]; + let mut z = vec![0.0; n]; + let mut d = vec![0.0; n]; + let mut q = vec![0.0; n]; + let true_residual = + |p: &[f64], r: &mut [f64], fine: &Level| -> f64 { fine.residual(&b, p, r) }; + let anchor = anchor.filter(|&a| a < n && fine.active[a]); + let t_iter = std::time::Instant::now(); + let finish = |p: &mut [f64], iterations: usize, residual: f64| { + for (c, members) in components.members.iter().enumerate() { + if !components.singular[c] { + continue; + } + let shift = match anchor { + Some(a) if components.id[a] == c => p[a], + _ => members.iter().map(|&idx| p[idx]).sum::() / members.len() as f64, + }; + for &idx in members { + p[idx] -= shift; + } + } + PoissonSolution { + iterations, + residual, + converged: residual < tolerance, + setup_ns, + iterate_ns: t_iter.elapsed().as_nanos() as u64, + } + }; + + let mut res = true_residual(p, &mut r, fine); + if res < tolerance { + return finish(p, 0, res); + } + precond(&r, &mut z); + if singular { + project_mean(&mut z); + } + for &idx in cells { + d[idx] = z[idx]; + } + let mut rz = dot(&r, &z); + let mut iterations = 0; + let mut last_true = res; + while iterations < params.max_iterations { + iterations += 1; + fine.apply(&d, &mut q); + let dq = dot(&d, &q); + if !dq.is_finite() || dq <= 0.0 || !rz.is_finite() || rz <= 0.0 { + res = true_residual(p, &mut r, fine); + return finish(p, iterations, res); + } + let alpha = rz / dq; + for &idx in cells { + p[idx] += alpha * d[idx]; + r[idx] -= alpha * q[idx]; + } + if l1(&r) < tolerance { + res = true_residual(p, &mut r, fine); + if res < tolerance || res > 0.9 * last_true { + return finish(p, iterations, res); + } + last_true = res; + } + precond(&r, &mut z); + if singular { + project_mean(&mut z); + } + let rz_new = dot(&r, &z); + let beta = rz_new / rz; + rz = rz_new; + for &idx in cells { + d[idx] = z[idx] + beta * d[idx]; + } + } + res = true_residual(p, &mut r, fine); + finish(p, iterations, res) +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/problem.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/problem.rs new file mode 100644 index 0000000..749301c --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/problem.rs @@ -0,0 +1,256 @@ +//! The seven-point problem `ap p − Σ a_nb p_nb = rhs` on the active cells, +//! `ap = ae + aw + an + as + at + ab + extra_diag`. The z direction may be +//! periodic; the neighbour above/below a cell is computed HERE, once, and +//! stored as data by the hierarchy — no stencil branches on it. + +#[derive(Debug, Clone)] +pub struct Problem { + pub nx: usize, + pub ny: usize, + pub nz: usize, + /// `k = nz − 1` neighbours `k = 0`. + pub periodic_z: bool, + pub active: Vec, + pub ae: Vec, + pub aw: Vec, + pub an: Vec, + pub as_: Vec, + /// Towards `k + 1`. + pub at: Vec, + /// Towards `k − 1`. + pub ab: Vec, + pub extra_diag: Vec, + pub rhs: Vec, +} + +impl Problem { + /// All cells active, every coefficient and the right-hand side zero. + #[must_use] + pub fn new(nx: usize, ny: usize, nz: usize) -> Self { + let n = nx * ny * nz; + Self { + nx, + ny, + nz, + periodic_z: false, + active: vec![true; n], + ae: vec![0.0; n], + aw: vec![0.0; n], + an: vec![0.0; n], + as_: vec![0.0; n], + at: vec![0.0; n], + ab: vec![0.0; n], + extra_diag: vec![0.0; n], + rhs: vec![0.0; n], + } + } + + #[inline] + #[must_use] + pub fn index(&self, k: usize, j: usize, i: usize) -> usize { + (k * self.ny + j) * self.nx + i + } + + #[inline] + #[must_use] + pub fn nxy(&self) -> usize { + self.nx * self.ny + } + + /// The cell above `idx` (plane `k`), wrapping when periodic. + #[inline] + #[must_use] + pub fn top(&self, idx: usize, k: usize) -> Option { + if k + 1 < self.nz { + Some(idx + self.nxy()) + } else if self.periodic_z && self.nz > 1 { + Some(idx - (self.nz - 1) * self.nxy()) + } else { + None + } + } + + #[inline] + #[must_use] + pub fn bottom(&self, idx: usize, k: usize) -> Option { + if k > 0 { + Some(idx - self.nxy()) + } else if self.periodic_z && self.nz > 1 { + Some(idx + (self.nz - 1) * self.nxy()) + } else { + None + } + } + + /// Diagonal `ap` (the 2D sum with the z terms appended). + #[inline] + #[must_use] + pub fn diagonal(&self, idx: usize) -> f64 { + self.ae[idx] + + self.aw[idx] + + self.an[idx] + + self.as_[idx] + + self.at[idx] + + self.ab[idx] + + self.extra_diag[idx] + } + + /// Pure Neumann: no active cell has a Dirichlet contribution. + #[must_use] + pub fn is_singular(&self) -> bool { + !self + .active + .iter() + .zip(&self.extra_diag) + .any(|(&a, &d)| a && d > 0.0) + } + + /// `Σ |rhs − (ap p − Σ a_nb p_nb)|` over the active cells with `ap > 0`. + #[must_use] + pub fn residual_l1(&self, p: &[f64]) -> f64 { + let (nx, ny, nz) = (self.nx, self.ny, self.nz); + let mut sum = 0.0; + for k in 0..nz { + for j in 0..ny { + for i in 0..nx { + let idx = self.index(k, j, i); + if !self.active[idx] { + continue; + } + let ap = self.diagonal(idx); + if ap <= 0.0 { + continue; + } + let mut nb = 0.0; + if i + 1 < nx && self.active[idx + 1] { + nb += self.ae[idx] * p[idx + 1]; + } + if i > 0 && self.active[idx - 1] { + nb += self.aw[idx] * p[idx - 1]; + } + if j + 1 < ny && self.active[idx + nx] { + nb += self.an[idx] * p[idx + nx]; + } + if j > 0 && self.active[idx - nx] { + nb += self.as_[idx] * p[idx - nx]; + } + if let Some(t) = self.top(idx, k) { + if self.active[t] { + nb += self.at[idx] * p[t]; + } + } + if let Some(b) = self.bottom(idx, k) { + if self.active[b] { + nb += self.ab[idx] * p[b]; + } + } + sum += (self.rhs[idx] - (ap * p[idx] - nb)).abs(); + } + } + } + sum + } + + /// Lengths, non-negativity, zero coefficients across domain edges and + /// towards inactive cells, symmetry to `1e-12` relative in x, y and z. + pub fn validate(&self) -> Result<(), String> { + let n = self.nx * self.ny * self.nz; + for (name, len) in [ + ("active", self.active.len()), + ("ae", self.ae.len()), + ("aw", self.aw.len()), + ("an", self.an.len()), + ("as_", self.as_.len()), + ("at", self.at.len()), + ("ab", self.ab.len()), + ("extra_diag", self.extra_diag.len()), + ("rhs", self.rhs.len()), + ] { + if len != n { + return Err(format!("{name}: length {len}, expected nx*ny*nz = {n}")); + } + } + let (nx, ny, nz) = (self.nx, self.ny, self.nz); + let symmetric = |a: f64, b: f64| (a - b).abs() <= 1e-12 * a.abs().max(b.abs()); + for k in 0..nz { + for j in 0..ny { + for i in 0..nx { + let idx = self.index(k, j, i); + for (name, v) in [ + ("ae", self.ae[idx]), + ("aw", self.aw[idx]), + ("an", self.an[idx]), + ("as_", self.as_[idx]), + ("at", self.at[idx]), + ("ab", self.ab[idx]), + ("extra_diag", self.extra_diag[idx]), + ] { + if v.is_nan() || v < 0.0 { + return Err(format!("{name}[{idx}] = {v} is negative or NaN")); + } + } + let active = self.active[idx]; + let check = |name: &str, coef: f64, nb_ok: bool| -> Result<(), String> { + if !nb_ok && coef != 0.0 { + return Err(format!( + "{name}[{idx}] = {coef} across a domain edge or towards an inactive cell (cell active = {active})" + )); + } + Ok(()) + }; + check( + "ae", + self.ae[idx], + active && i + 1 < nx && self.active[idx + 1], + )?; + check("aw", self.aw[idx], active && i > 0 && self.active[idx - 1])?; + check( + "an", + self.an[idx], + active && j + 1 < ny && self.active[idx + nx], + )?; + check( + "as_", + self.as_[idx], + active && j > 0 && self.active[idx - nx], + )?; + let t = self.top(idx, k); + let b = self.bottom(idx, k); + check( + "at", + self.at[idx], + active && t.is_some_and(|t| self.active[t]), + )?; + check( + "ab", + self.ab[idx], + active && b.is_some_and(|b| self.active[b]), + )?; + if i + 1 < nx && !symmetric(self.ae[idx], self.aw[idx + 1]) { + return Err(format!( + "asymmetric x face at {idx}: ae {} vs aw {}", + self.ae[idx], + self.aw[idx + 1] + )); + } + if j + 1 < ny && !symmetric(self.an[idx], self.as_[idx + nx]) { + return Err(format!( + "asymmetric y face at {idx}: an {} vs as {}", + self.an[idx], + self.as_[idx + nx] + )); + } + if let Some(t) = t { + if !symmetric(self.at[idx], self.ab[t]) { + return Err(format!( + "asymmetric z face at {idx}: at {} vs ab {}", + self.at[idx], self.ab[t] + )); + } + } + } + } + } + Ok(()) + } +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs index facf67a..d899647 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs @@ -17,6 +17,7 @@ pub mod boundary_conditions; pub mod curvilinear; /// PISO on the fixed grid with an embedded body pub mod embedded; +pub mod embedded3; /// Embedded-body geometry, classification and loads pub mod embedded_body; /// Flow field data structures diff --git a/crates/specialized/rtx-cfd/tests/embedded3_poisson_identity.rs b/crates/specialized/rtx-cfd/tests/embedded3_poisson_identity.rs new file mode 100644 index 0000000..0f88c67 --- /dev/null +++ b/crates/specialized/rtx-cfd/tests/embedded3_poisson_identity.rs @@ -0,0 +1,222 @@ +//! embedded3 gate 1 (omni-cortex `docs/embedded3_campaign.md`): the Poisson +//! solver at `nz = 1` is the 2D solver bit for bit (and the `three_d` +//! oracle's); an extrusion in z agrees across planes to the solve's +//! accuracy, bit-identically for lexicographic decoupled planes. + +use rtx_cfd::solvers::incompressible::embedded3::poisson::{ + PcgCache, Problem, solve_pcg, solve_pcg_cached, +}; +use rtx_cfd::solvers::incompressible::three_d::poisson::{PoissonProblem3D, solve_multigrid_pcg3}; +use rtx_cfd::solvers::incompressible::{ + MgSmoother, MultigridParameters, PcgCache as PcgCache2, PoissonProblem, solve_multigrid_pcg, + solve_multigrid_pcg_cached, +}; + +/// The `poisson_redblack.rs` masked channel. +fn problem_2d(nx: usize, ny: usize, seed: u64) -> PoissonProblem { + let mut p = PoissonProblem::new(nx, ny); + let (dx, dy, dt) = (1.0 / nx as f64, 0.41 / ny as f64, 1e-3); + let (ae, an) = (dt * dy / dx, dt * dx / dy); + let hole = |i: usize, j: usize| { + let (x, y) = ((i as f64 + 0.5) * dx, (j as f64 + 0.5) * dy); + (x - 0.2).powi(2) + (y - 0.2).powi(2) < 0.05 * 0.05 + }; + for j in 0..ny { + for i in 0..nx { + let idx = j * nx + i; + if hole(i, j) { + p.active[idx] = false; + continue; + } + if i + 1 < nx && !hole(i + 1, j) { + p.ae[idx] = ae; + } + if i > 0 && !hole(i - 1, j) { + p.aw[idx] = ae; + } + if j + 1 < ny && !hole(i, j + 1) { + p.an[idx] = an; + } + if j > 0 && !hole(i, j - 1) { + p.as_[idx] = an; + } + if i + 1 == nx { + p.extra_diag[idx] = 2.0 * ae; + } + } + } + let mut state = seed | 1; + for idx in 0..nx * ny { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + p.rhs[idx] = if p.active[idx] { + 1e-6 * ((state >> 11) as f64 / (1u64 << 53) as f64 - 0.5) + } else { + 0.0 + }; + } + p +} + +/// The 2D problem stacked `nz` times; `az` couples the planes. +fn extrude(p2: &PoissonProblem, nz: usize, az: f64, periodic_z: bool) -> Problem { + let (nx, ny) = (p2.nx, p2.ny); + let mut p = Problem::new(nx, ny, nz); + p.periodic_z = periodic_z; + for k in 0..nz { + for idx2 in 0..nx * ny { + let idx = k * nx * ny + idx2; + p.active[idx] = p2.active[idx2]; + p.ae[idx] = p2.ae[idx2]; + p.aw[idx] = p2.aw[idx2]; + p.an[idx] = p2.an[idx2]; + p.as_[idx] = p2.as_[idx2]; + p.extra_diag[idx] = p2.extra_diag[idx2]; + p.rhs[idx] = p2.rhs[idx2]; + if p2.active[idx2] && az != 0.0 { + if k + 1 < nz || periodic_z { + p.at[idx] = az; + } + if k > 0 || periodic_z { + p.ab[idx] = az; + } + } + } + } + p +} + +fn to_three_d(p: &Problem) -> PoissonProblem3D { + let mut q = PoissonProblem3D::new(p.nx, p.ny, p.nz); + q.periodic_z = p.periodic_z; + q.active.clone_from(&p.active); + q.ae.clone_from(&p.ae); + q.aw.clone_from(&p.aw); + q.an.clone_from(&p.an); + q.as_.clone_from(&p.as_); + q.at.clone_from(&p.at); + q.ab.clone_from(&p.ab); + q.extra_diag.clone_from(&p.extra_diag); + q.rhs.clone_from(&p.rhs); + q +} + +fn bits(v: &[f64]) -> Vec { + v.iter().map(|x| x.to_bits()).collect() +} + +#[test] +fn nz_one_is_the_two_d_solver_bit_for_bit() { + let (nx, ny) = (96, 40); + let tol = 1e-12; + for (name, params) in [ + ("lexicographic", MultigridParameters::default()), + ( + "red-black", + MultigridParameters { + smoother: MgSmoother::RedBlack, + ..MultigridParameters::default() + }, + ), + ] { + let mut cache2 = PcgCache2::default(); + let mut cache3 = PcgCache::default(); + for seed in [5u64, 20, 21] { + let p2 = problem_2d(nx, ny, seed); + let p3 = extrude(&p2, 1, 0.0, false); + assert!(p3.validate().is_ok(), "{:?}", p3.validate()); + let (mut a, mut b) = (vec![0.0; nx * ny], vec![0.0; nx * ny]); + let sa = solve_multigrid_pcg(&p2, &mut a, ¶ms, tol, None); + let sb = solve_pcg(&p3, &mut b, ¶ms, tol, None); + assert!(sa.converged && sb.converged); + assert_eq!(sa.iterations, sb.iterations, "{name} seed {seed}"); + assert_eq!( + bits(&a), + bits(&b), + "{name} seed {seed}: embedded3 differs from the 2D solver" + ); + let (mut c, mut d) = (vec![0.0; nx * ny], vec![0.0; nx * ny]); + let sc = solve_multigrid_pcg_cached(&p2, &mut c, ¶ms, tol, None, &mut cache2); + let sd = solve_pcg_cached(&p3, &mut d, ¶ms, tol, None, &mut cache3); + assert_eq!(sc.iterations, sd.iterations); + assert_eq!(bits(&c), bits(&d), "{name} seed {seed}: cached differs"); + assert_eq!(bits(&a), bits(&c)); + assert!(p3.residual_l1(&b) < tol); + println!( + " {name} seed {seed}: {} iterations, bit-identical to the 2D solver", + sa.iterations + ); + } + } +} + +/// Bit identity across planes for lexicographic decoupled planes; the +/// solve's accuracy (`1e-6·scale`) elsewhere (the red-black colouring swaps +/// between planes; coupled lexicographic sweeps read new values below). +/// Every case is also bit-identical to the `three_d` oracle. +#[test] +fn an_extrusion_in_z_is_z_invariant_and_equals_the_oracle() { + let (nx, ny, nz) = (48, 20, 8); + let tol = 1e-12; + let p2 = problem_2d(nx, ny, 7); + let az = 1e-3 * (1.0 / 48.0) * (0.41 / 20.0) / 0.05; + for (name, az, periodic, decoupled) in [ + ("decoupled planes", 0.0, false, true), + ("periodic z", az, true, false), + ("closed z (walls)", az, false, false), + ] { + for smoother in [MgSmoother::Lexicographic, MgSmoother::RedBlack] { + let params = MultigridParameters { + smoother, + ..MultigridParameters::default() + }; + let p3 = extrude(&p2, nz, az, periodic); + assert!(p3.validate().is_ok(), "{name}: {:?}", p3.validate()); + let mut sol = vec![0.0; nx * ny * nz]; + let s = solve_pcg(&p3, &mut sol, ¶ms, tol, None); + assert!( + s.converged, + "{name} {smoother:?}: not converged ({} it)", + s.iterations + ); + assert!(p3.residual_l1(&sol) < tol); + let mut oracle = vec![0.0; nx * ny * nz]; + let so = solve_multigrid_pcg3(&to_three_d(&p3), &mut oracle, ¶ms, tol, None); + assert_eq!( + so.iterations, s.iterations, + "{name} {smoother:?}: oracle iterations" + ); + assert_eq!( + bits(&oracle), + bits(&sol), + "{name} {smoother:?}: differs from the three_d oracle" + ); + let plane = |k: usize| &sol[k * nx * ny..(k + 1) * nx * ny]; + let scale = sol.iter().fold(0.0_f64, |m, v| m.max(v.abs())); + let mut worst = 0.0_f64; + for k in 1..nz { + let d = plane(k) + .iter() + .zip(plane(0)) + .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs())); + worst = worst.max(d); + if decoupled && smoother == MgSmoother::Lexicographic { + assert_eq!( + bits(plane(k)), + bits(plane(0)), + "{name}: plane {k} differs in bits" + ); + } + } + assert!( + worst <= 1e-6 * scale, + "{name} {smoother:?}: planes differ by {worst:.3e} of {scale:.3e}" + ); + println!( + " {name} {smoother:?}: {} iterations, planes within {worst:.2e} of {scale:.2e}, = three_d oracle", + s.iterations + ); + } + } +}