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

//! PERF-2 (`docs/perf2_campaign.md`): the red-black symmetric Gauss–Seidel
//! smoother is a different preconditioner, not a bit-identical one: the pin
//! is that it solves the same masked problem to the same stop, that its
//! solution agrees with the lexicographic one to the solver's tolerance,
//! and that the cached red-black solve is bit-identical to the uncached
//! red-black solve (the cache keys on the smoother).

use rtx_cfd::solvers::incompressible::{
    MgSmoother, MultigridParameters, PcgCache, PoissonProblem, solve_multigrid_pcg,
    solve_multigrid_pcg_cached,
};

fn problem(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
}

#[test]
fn red_black_solves_the_masked_problem_and_caches_exactly() {
    let (nx, ny) = (96, 40);
    let prob = problem(nx, ny, 5);
    let tol = 1e-12;
    let lex = MultigridParameters::default();
    let rb = MultigridParameters {
        smoother: MgSmoother::RedBlack,
        ..MultigridParameters::default()
    };
    let (mut p_lex, mut p_rb) = (vec![0.0; nx * ny], vec![0.0; nx * ny]);
    let s_lex = solve_multigrid_pcg(&prob, &mut p_lex, &lex, tol, None);
    let s_rb = solve_multigrid_pcg(&prob, &mut p_rb, &rb, tol, None);
    assert!(s_lex.converged && s_rb.converged);
    assert!(
        prob.residual_l1(&p_rb) < tol,
        "red-black residual {:.3e}",
        prob.residual_l1(&p_rb)
    );
    let scale = p_lex.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
    let diff = p_lex
        .iter()
        .zip(&p_rb)
        .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
    println!(
        "  lexicographic {} iterations vs red-black {} iterations; solutions differ by {:.3e} on a scale of {:.3e}",
        s_lex.iterations, s_rb.iterations, diff, scale
    );
    assert!(
        diff < 1e-6 * scale,
        "red-black and lexicographic disagree: {diff:.3e} of {scale:.3e}"
    );
    // The cached red-black solve is the uncached one, bit for bit.
    let mut cache = PcgCache::default();
    for k in 0..3u64 {
        let mut q = prob.clone();
        q.rhs = problem(nx, ny, 20 + k).rhs;
        let (mut a, mut b) = (vec![0.0; nx * ny], vec![0.0; nx * ny]);
        let sa = solve_multigrid_pcg(&q, &mut a, &rb, tol, None);
        let sb = solve_multigrid_pcg_cached(&q, &mut b, &rb, tol, None, &mut cache);
        assert!(a.iter().zip(&b).all(|(x, y)| x.to_bits() == y.to_bits()));
        assert_eq!(sa.iterations, sb.iterations);
    }
    // Switching the smoother is a cache miss (the key carries it), still exact.
    let (mut a, mut b) = (vec![0.0; nx * ny], vec![0.0; nx * ny]);
    let sa = solve_multigrid_pcg(&prob, &mut a, &lex, tol, None);
    let sb = solve_multigrid_pcg_cached(&prob, &mut b, &lex, tol, None, &mut cache);
    assert!(a.iter().zip(&b).all(|(x, y)| x.to_bits() == y.to_bits()));
    assert_eq!(sa.iterations, sb.iterations);
}