rtx-cfd 3D Stage 1 item 1: three_d::{Grid3, poisson} — the 2D Poisson stack transcribed to a seven-point operator (sanitised coefficients, 2×2×2 Galerkin aggregation, (i+j+k)%2 colouring, periodic z, run_pcg line for line, PcgCache3); gate 1 HELD: nz=1 bit-identical to the 2D solver (solution + iterations, lex + red-black, cached/uncached); extrusion z-invariant to the solve's accuracy (bit-identical planes for lexicographic decoupled); probes for the aggregation/colouring interaction
CI / Test (macos-latest) (push) Blocked by required conditions
CI / Test (ubuntu-latest) (push) Blocked by required conditions
CI / Python Bindings (maturin) (macos-latest) (push) Blocked by required conditions
CI / Python Bindings (maturin) (ubuntu-latest) (push) Blocked by required conditions
CI / WASM Build + Size Check (push) Blocked by required conditions
CI / Distributed Training Tests (push) Blocked by required conditions
CI / CI Success (push) Blocked by required conditions
CI / Build (macos-latest) (push) Waiting to run
Documentation / Build API Documentation (push) Failing after 4s
CI / Build (ubuntu-latest) (push) Failing after 4s
Documentation / Build User Guide (push) Successful in 5s
CI / Format Check (push) Failing after 10s
CI / Build CPU-Only (Explicit) (push) Failing after 1m38s
CI / Clippy Check (push) Failing after 2m43s
Performance Benchmarks / Run Benchmarks (push) Successful in 6m24s

Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
Omar Sobh
2026-09-17 11:11:52 -05:00
co-authored by Claude Fable 5.1
parent 7e135893a2
commit 0794e1b6db
4 changed files with 1376 additions and 0 deletions
@@ -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::{
@@ -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
}
}
File diff suppressed because it is too large Load Diff
@@ -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<u64> {
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, &params, tol, None);
let sb = solve_multigrid_pcg3(&p3, &mut b, &params, 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, &params, tol, None, &mut cache2);
let sd = solve_multigrid_pcg3_cached(&p3, &mut d, &params, 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, &params, 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, &params, 1e-12, None);
println!(
" nz {nz} az/ae {:.0} {smoother:?} coarsest {}: {} iterations",
az / ae,
if coarsest == 32 { "32" } else { "single level" },
s.iterations
);
}
}
}
}