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

Co-Authored-By: Claude Fable 5.1 <[email protected]>
2026-09-17 11:11:52 -05:00
parent 7e135893a2
commit 0794e1b6db
4 changed files with 1376 additions and 0 deletions
@@ -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
+                );
+            }
+        }
+    }
+}