rustytorch/crates/specialized/rtx-cfd/tests/embedded3_poisson_identity.rs

//! 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<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 = 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, &params, tol, None);
            let sb = solve_pcg(&p3, &mut b, &params, 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, &params, tol, None, &mut cache2);
            let sd = solve_pcg_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 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, &params, 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, &params, 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
            );
        }
    }
}