diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs index 0074de5..facf67a 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/mod.rs @@ -38,6 +38,7 @@ pub mod simple; pub mod simple_gpu; /// CSR matrix + Jacobi-BiCGSTAB for the curvilinear pressure equation pub mod sparse_bicgstab; +pub mod three_d; // Re-export main types pub use ale::{ diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/mod.rs new file mode 100644 index 0000000..596541e --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/mod.rs @@ -0,0 +1,67 @@ +//! The three-dimensional embedded solver (omni-cortex +//! `docs/three_d_stage1_campaign.md`): a sharp-interface embedded wall on a +//! Cartesian grid, device-resident. Stage 1 = the core (this module tree), +//! the smooth wall and the DFG 3D-2Z gate. The 2D solver is NOT touched: +//! at `nz = 1` the code here must reproduce its digits, which is the first +//! gate of every piece. +//! +//! Layout: cells are `(k, j, i)` row-major, `cell = (k * ny + j) * nx + i`. + +pub mod poisson; + +/// A uniform Cartesian grid. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct Grid3 { + pub nx: usize, + pub ny: usize, + pub nz: usize, + pub dx: f64, + pub dy: f64, + pub dz: f64, +} + +impl Grid3 { + /// Cells in the domain. + #[inline] + #[must_use] + pub fn cells(&self) -> usize { + self.nx * self.ny * self.nz + } + + /// Row-major cell index of `(k, j, i)`. + #[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) + } + + /// Index of the u face west of cell `(k, j, i)` on the `(nx + 1) × ny × nz` + /// staggered array (`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 + } + + /// Index of the v face south of cell `(k, j, i)` on `nx × (ny + 1) × nz`. + #[inline] + #[must_use] + pub fn vface(&self, k: usize, j: usize, i: usize) -> usize { + (k * (self.ny + 1) + j) * self.nx + i + } + + /// Index of the w face below cell `(k, j, i)` on `nx × ny × (nz + 1)`. + #[inline] + #[must_use] + pub fn wface(&self, k: usize, j: usize, i: usize) -> usize { + (k * self.ny + j) * self.nx + i + } +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/poisson/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/poisson/mod.rs new file mode 100644 index 0000000..0075f06 --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/poisson/mod.rs @@ -0,0 +1,1069 @@ +//! The 3D pressure Poisson solver: the 2D `poisson.rs` transcribed to a +//! seven-point operator on `(k, j, i)` cells, keeping every rule of the 2D +//! code (sanitised coefficients, Galerkin aggregation by 2 per direction, +//! the ×2 coarse correction, red-black colouring by `(i + j + k) % 2`, the +//! f64 CG with the true-residual stop). At `nz = 1` with zero z-coefficients +//! the arithmetic is the 2D solver's in the same order (gate 1: bit +//! identity). `periodic_z` closes the z direction (the z-invariance oracle +//! of later gates). + +use crate::solvers::incompressible::poisson::{ + MgPrecision, MgScalar, MgSmoother, MultigridParameters, PoissonSolution, +}; + +/// Symmetric GS sweeps on the coarsest level (the 2D value). +const COARSEST_SWEEPS: usize = 50; +/// The 2D `COARSE_CORRECTION`, proved dimension-independent there. +const COARSE_CORRECTION: f64 = 2.0; +const MAX_LEVELS: usize = 64; + +/// The seven-point problem: `ap p − Σ a_nb p_nb = rhs` on the active cells, +/// `ap = ae + aw + an + as + at + ab + extra_diag`. +#[derive(Debug, Clone)] +pub struct PoissonProblem3D { + pub nx: usize, + pub ny: usize, + pub nz: usize, + /// The z direction wraps (`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 PoissonProblem3D { + /// 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 `(k)` (wrapping when periodic), or `None`. + #[inline] + 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] + 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` of cell `idx` (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: ae[{idx}] = {} vs aw[{}] = {}", + self.ae[idx], + idx + 1, + self.aw[idx + 1] + )); + } + if j + 1 < ny && !symmetric(self.an[idx], self.as_[idx + nx]) { + return Err(format!( + "asymmetric y face: an[{idx}] = {} vs as_[{}] = {}", + self.an[idx], + idx + nx, + 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}] = {} vs ab[{t}] = {}", + self.at[idx], self.ab[t] + )); + } + } + } + } + } + Ok(()) + } +} + +/// One level: the problem, its coefficients in the V-cycle scalar, the +/// active cells, the two colours, the parent map. +#[derive(Clone)] +struct Level3 { + problem: PoissonProblem3D, + active: Vec, + ae: Vec, + aw: Vec, + an: Vec, + as_: Vec, + at: Vec, + ab: Vec, + ap: Vec, + cells: Vec, + red: Vec, + black: Vec, + /// Neighbour above / below per cell (`usize::MAX` = none); explicit so + /// the periodic wrap costs nothing in the stencil. + top: Vec, + bot: Vec, + coarse_of: Vec, +} + +struct Work3 { + b: Vec, + x: Vec, + r: Vec, +} + +impl Work3 { + fn new(n: usize) -> Self { + Self { + b: vec![T::ZERO; n], + x: vec![T::ZERO; n], + r: vec![T::ZERO; n], + } + } +} + +impl Level3 { + fn new(mut problem: PoissonProblem3D) -> 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 + } + + 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); + } + } + + 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 &idx in &self.red { + x[idx] = (b[idx] + self.neighbour_sum(x, idx)) / self.ap[idx]; + } + for &idx in &self.black { + x[idx] = (b[idx] + self.neighbour_sum(x, idx)) / self.ap[idx]; + } + for &idx in &self.black { + x[idx] = (b[idx] + self.neighbour_sum(x, idx)) / self.ap[idx]; + } + for &idx in &self.red { + 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) -> (PoissonProblem3D, 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 = PoissonProblem3D::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 { + let kt = self.top[idx] / (nx * ny); + if cz(kt) != kc { + coarse.at[c] += p.at[idx]; + } + } + if p.ab[idx] != 0.0 { + let kb = self.bot[idx] / (nx * ny); + if cz(kb) != kc { + coarse.ab[c] += p.ab[idx]; + } + } + } + (coarse, coarse_of) + } +} + +/// The hierarchy: level 0 is the fine problem. +pub(crate) struct Hierarchy3 { + levels: Vec>, + work: Vec>, + sweeps: usize, + smoother: MgSmoother, +} + +impl Hierarchy3 { + pub(crate) fn build(problem: &PoissonProblem3D, params: &MultigridParameters) -> Self { + let mut levels = vec![Level3::::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 = Level3::::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| Work3::new(l.problem.nx * l.problem.ny * l.problem.nz)) + .collect(); + Self { + levels, + work, + sweeps: params.smoother_sweeps.max(1), + smoother: params.smoother, + } + } + + pub(crate) fn depth(&self) -> usize { + self.levels.len() + } + + pub(crate) fn cells(&self, l: usize) -> &[usize] { + &self.levels[l].cells + } + + /// `z = M⁻¹ r`: one V-cycle from zero (the 2D sequence). + pub(crate) 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 Work3 { 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 Work3 { 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 Work3 { 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(); + } + } +} + +/// The operator part of a problem plus the hierarchy parameters (bit +/// patterns), to decide reuse of a prepared solver. +struct OperatorKey3 { + nx: usize, + ny: usize, + nz: usize, + periodic_z: bool, + active: Vec, + coefficients: Vec, + smoother_sweeps: usize, + coarsest_cells: usize, + smoother: MgSmoother, +} + +impl OperatorKey3 { + fn bits(problem: &PoissonProblem3D) -> 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()) + } + + fn of(problem: &PoissonProblem3D, 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, + } + } + + fn matches(&self, problem: &PoissonProblem3D, 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)) + } +} + +/// Connected components of the active cells through non-zero faces, and +/// whether each is singular. +#[derive(Clone)] +struct Components3 { + id: Vec, + members: Vec>, + singular: Vec, +} + +impl Components3 { + fn find(problem: &PoissonProblem3D, 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, + } + } +} + +/// Everything the CG derives from the operator. +struct Prepared3 { + key: OperatorKey3, + hier: Hierarchy3, + fine: Level3, + cells: Vec, + components: Components3, +} + +impl Prepared3 { + fn build(problem: &PoissonProblem3D, params: &MultigridParameters) -> Self { + let hier = Hierarchy3::::build(problem, params); + let fine = Level3::::new(problem.clone()); + let cells: Vec = fine.cells.clone(); + let components = Components3::find(problem, &cells); + Self { + key: OperatorKey3::of(problem, params), + hier, + fine, + cells, + components, + } + } +} + +/// Prepared operators, reused when the operator is bit-identical. +#[derive(Default)] +pub struct PcgCache3 { + f64: Option>, + f32: Option>, +} + +/// [`solve_multigrid_pcg3`] with the operator taken from `cache` on a hit. +pub fn solve_multigrid_pcg3_cached( + problem: &PoissonProblem3D, + p: &mut [f64], + params: &MultigridParameters, + tolerance: f64, + anchor: Option, + cache: &mut PcgCache3, +) -> 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) + } + } +} + +/// CG on the active cells preconditioned by one V-cycle (the 2D solver's +/// contract: absolute L1 true-residual stop, per-component mean projection +/// on singular components, anchor shift on exit). +pub fn solve_multigrid_pcg3( + problem: &PoissonProblem3D, + p: &mut [f64], + params: &MultigridParameters, + tolerance: f64, + anchor: Option, +) -> PoissonSolution { + match params.precision { + MgPrecision::F64 => solve_pcg_with::(problem, p, params, tolerance, anchor), + MgPrecision::F32 => solve_pcg_with::(problem, p, params, tolerance, anchor), + } +} + +fn solve_cached_with( + problem: &PoissonProblem3D, + 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(Prepared3::::build(problem, params)); + } + let setup_ns = t_entry.elapsed().as_nanos() as u64; + let prep = slot.as_mut().expect("prepared"); + run_pcg3(prep, problem, p, params, tolerance, anchor, setup_ns) +} + +fn solve_pcg_with( + problem: &PoissonProblem3D, + p: &mut [f64], + params: &MultigridParameters, + tolerance: f64, + anchor: Option, +) -> PoissonSolution { + let t_entry = std::time::Instant::now(); + let mut prep = Prepared3::::build(problem, params); + let setup_ns = t_entry.elapsed().as_nanos() as u64; + run_pcg3(&mut prep, problem, p, params, tolerance, anchor, setup_ns) +} + +/// The 2D `run_pcg`, line for line. +fn run_pcg3( + prep: &mut Prepared3, + problem: &PoissonProblem3D, + 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 PoissonProblem3D: {:?}", + problem.validate() + ); + let Prepared3 { + hier, + fine, + cells, + components, + .. + } = prep; + let fine: &Level3 = fine; + let cells: &[usize] = cells; + let components: &Components3 = components; + let mut precond = |r: &[f64], z: &mut [f64]| { + hier.apply_preconditioner(r, z); + }; + let active_n = cells.len(); + if active_n == 0 { + return PoissonSolution { + iterations: 0, + residual: 0.0, + converged: true, + setup_ns, + iterate_ns: 0, + }; + } + let singular_any = 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 singular = singular_any; + + 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: &Level3| -> 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) +} + +#[cfg(test)] +mod probe { + use super::*; + + fn stacked(nx: usize, ny: usize, nz: usize) -> PoissonProblem3D { + let mut p = PoissonProblem3D::new(nx, ny, nz); + let (ae, an) = (1.0e-3, 0.98e-3); + for k in 0..nz { + for j in 0..ny { + for i in 0..nx { + let idx = p.index(k, j, i); + if i + 1 < nx { + p.ae[idx] = ae; + } + if i > 0 { + p.aw[idx] = ae; + } + if j + 1 < ny { + p.an[idx] = an; + } + if j > 0 { + p.as_[idx] = an; + } + if i + 1 == nx { + p.extra_diag[idx] = 2.0 * ae; + } + } + } + } + p + } + + fn symmetry_defect( + p: &PoissonProblem3D, + smoother: MgSmoother, + coarsest: usize, + ) -> (usize, f64) { + let params = MultigridParameters { + smoother, + coarsest_cells: coarsest, + ..MultigridParameters::default() + }; + let mut h = Hierarchy3::::build(p, ¶ms); + let n = p.nx * p.ny * p.nz; + let mut state = 12345u64; + let mut rnd = || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state >> 11) as f64 / (1u64 << 53) as f64 - 0.5 + }; + let x: Vec = (0..n).map(|_| rnd()).collect(); + let y: Vec = (0..n).map(|_| rnd()).collect(); + let (mut mx, mut my) = (vec![0.0; n], vec![0.0; n]); + h.apply_preconditioner(&x, &mut mx); + h.apply_preconditioner(&y, &mut my); + let dot = |a: &[f64], b: &[f64]| a.iter().zip(b).map(|(u, v)| u * v).sum::(); + let (a, b) = (dot(&x, &my), dot(&mx, &y)); + (h.depth(), (a - b).abs() / a.abs().max(b.abs())) + } + + #[test] + #[ignore] + fn probe_symmetry() { + for nz in [1usize, 2, 4] { + for smoother in [MgSmoother::Lexicographic, MgSmoother::RedBlack] { + for coarsest in [32usize, usize::MAX / 2] { + let p = stacked(48, 20, nz); + let (depth, d) = symmetry_defect(&p, smoother, coarsest); + println!(" nz {nz} {smoother:?} depth {depth}: symmetry defect {d:.3e}"); + } + } + } + } +} diff --git a/crates/specialized/rtx-cfd/tests/three_d_poisson_identity.rs b/crates/specialized/rtx-cfd/tests/three_d_poisson_identity.rs new file mode 100644 index 0000000..473b0d8 --- /dev/null +++ b/crates/specialized/rtx-cfd/tests/three_d_poisson_identity.rs @@ -0,0 +1,239 @@ +//! 3D Stage 1, gate 1 (omni-cortex `docs/three_d_stage1_campaign.md`): the +//! 3D Poisson solver at `nz = 1` is the 2D solver bit for bit (solution and +//! iteration count; lexicographic and red-black; cached and uncached), and +//! an extrusion in z (decoupled planes, and periodic z with a z-invariant +//! right-hand side) is bit-identical across planes. + +use rtx_cfd::solvers::incompressible::three_d::poisson::{ + PcgCache3, PoissonProblem3D, solve_multigrid_pcg3, solve_multigrid_pcg3_cached, +}; +use rtx_cfd::solvers::incompressible::{ + MgSmoother, MultigridParameters, PcgCache, PoissonProblem, solve_multigrid_pcg, + solve_multigrid_pcg_cached, +}; + +/// The `poisson_redblack.rs` masked channel (a hole, an outlet Dirichlet). +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 (0 = decoupled). +fn extrude(p2: &PoissonProblem, nz: usize, az: f64, periodic_z: bool) -> PoissonProblem3D { + let (nx, ny) = (p2.nx, p2.ny); + let mut p = PoissonProblem3D::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 { + let up = k + 1 < nz || periodic_z; + let down = k > 0 || periodic_z; + if up { + p.at[idx] = az; + } + if down { + p.ab[idx] = az; + } + } + } + } + p +} + +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 = PcgCache::default(); + let mut cache3 = PcgCache3::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_multigrid_pcg3(&p3, &mut b, ¶ms, tol, None); + assert!(sa.converged && sb.converged, "{name} seed {seed} converged"); + assert_eq!( + sa.iterations, sb.iterations, + "{name} seed {seed} iterations" + ); + assert_eq!(bits(&a), bits(&b), "{name} seed {seed}: 3D differs from 2D"); + 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_multigrid_pcg3_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 3D differs"); + assert_eq!( + bits(&a), + bits(&c), + "{name} seed {seed}: cached 2D differs from uncached" + ); + assert!(p3.residual_l1(&b) < tol); + println!( + " {name} seed {seed}: {} iterations, bit-identical to the 2D solver", + sa.iterations + ); + } + } +} + +/// Plane-to-plane identity of a z-invariant solve. Bit identity across +/// planes holds only where every plane's arithmetic path is the same: +/// decoupled planes under the lexicographic smoother. The seven-point +/// red-black colouring `(i + j + k) % 2` swaps the colours between +/// neighbouring planes on every level (the 2×2×2 aggregation merges plane +/// pairs, so the coarse levels swap again), and a lexicographic sweep of +/// coupled planes reads updated values below and old values above; in +/// both cases the planes agree to the solve's own accuracy (`1e-6 · scale`, +/// the red-black-vs-lexicographic pin's standard), not in bits. +#[test] +fn an_extrusion_in_z_is_z_invariant() { + 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_multigrid_pcg3(&p3, &mut sol, ¶ms, tol, None); + assert!( + s.converged, + "{name} {smoother:?}: not converged ({} it)", + s.iterations + ); + assert!(p3.residual_l1(&sol) < tol, "{name} {smoother:?}: residual"); + 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} {smoother:?}: plane {k} differs from plane 0 in bits" + ); + } + } + assert!( + worst <= 1e-6 * scale, + "{name} {smoother:?}: planes differ by {worst:.3e} on a scale of {scale:.3e}" + ); + println!( + " {name} {smoother:?}: {} iterations, planes within {worst:.2e} of {scale:.2e}{}", + s.iterations, + if decoupled && smoother == MgSmoother::Lexicographic { + " (planes bit-identical)" + } else { + "" + } + ); + } + } +} + +#[test] +#[ignore] +fn probe_iteration_counts() { + let (nx, ny) = (48, 20); + let p2 = problem_2d(nx, ny, 7); + let ae = 1e-3 * (0.41 / 20.0) / (1.0 / 48.0); + for (nz, az) in [(1usize, 0.0), (8, 0.0), (8, ae), (16, ae)] { + for smoother in [MgSmoother::Lexicographic, MgSmoother::RedBlack] { + for coarsest in [32usize, usize::MAX / 2] { + let params = MultigridParameters { + smoother, + coarsest_cells: coarsest, + ..MultigridParameters::default() + }; + let p3 = extrude(&p2, nz, az, true); + let mut sol = vec![0.0; nx * ny * nz]; + let s = solve_multigrid_pcg3(&p3, &mut sol, ¶ms, 1e-12, None); + println!( + " nz {nz} az/ae {:.0} {smoother:?} coarsest {}: {} iterations", + az / ae, + if coarsest == 32 { "32" } else { "single level" }, + s.iterations + ); + } + } + } +}