Omar SobhandClaude Fable 5.1
0dc95a8de5
CI / Build (macos-latest) (push) Waiting to run
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
Performance Benchmarks / Run Benchmarks (push) Failing after 6s
CI / Build (ubuntu-latest) (push) Failing after 5s
Documentation / Build API Documentation (push) Failing after 6s
Documentation / Build User Guide (push) Successful in 7s
CI / Format Check (push) Failing after 12s
CI / Clippy Check (push) Failing after 36s
CI / Build CPU-Only (Explicit) (push) Failing after 1m37s
Co-Authored-By: Claude Fable 5.1 <[email protected]> Claude-Session: https://claude.ai/code/session_01YJPeT6WA2e7YvAnS875AHL
1400 lines
51 KiB
Rust
1400 lines
51 KiB
Rust
//! Cell-centred five-point Poisson problems on an active-cell mask, solved
|
||
//! by conjugate gradient preconditioned with one V-cycle of geometric
|
||
//! (aggregation) multigrid.
|
||
//!
|
||
//! The pressure-correction systems of the fixed-grid PISO (`piso.rs`) and
|
||
//! the embedded-body PISO (`embedded.rs`) are both instances of
|
||
//! [`PoissonProblem`]: a symmetric positive (semi-)definite M-matrix whose
|
||
//! off-diagonal coefficients are face conductances (`dt dy/dx`, `dt dx/dy`),
|
||
//! zero across every prescribed face, with an optional diagonal-only
|
||
//! Dirichlet contribution (`extra_diag`) on cells next to a pressure outlet.
|
||
//! Both solvers used point SOR with the true-residual stop; SOR contracts
|
||
//! the smooth modes by `1 - O(h)` per sweep, so the sweep count grows with
|
||
//! the mesh and the 256² projection of a long march dominates the step.
|
||
//! [`solve_multigrid_pcg`] is the mesh-independent replacement.
|
||
//!
|
||
//! Design:
|
||
//! - The hierarchy is a `Vec` of levels, each holding its own
|
||
//! [`PoissonProblem`]. Coarsening is by 2 in each direction (odd sizes:
|
||
//! the last coarse cell aggregates what remains), a coarse cell is active
|
||
//! iff any child is. Prolongation is injection of the coarse value into
|
||
//! every child, restriction is the sum over children, and the coarse
|
||
//! operator is the Galerkin product `R A P` — which for constant
|
||
//! prolongation is exactly "sum the fine face coefficients across each
|
||
//! coarse face, sum the `extra_diag` of the children"; internal fine faces
|
||
//! cancel. The coarse problems therefore have the same structure as the
|
||
//! fine one, and the symmetry of the coarse operator is inherited.
|
||
//! - The prolongated coarse correction is scaled by [`COARSE_CORRECTION`]
|
||
//! `= 2`: constant prolongation under-corrects the smooth modes by exactly
|
||
//! that factor (see the constant's documentation), and without the scaling
|
||
//! the CG iteration count grows by ~10 per level.
|
||
//! - Singular blocks are handled per connected component of the active
|
||
//! cells: a component without Dirichlet data has its right-hand-side mean
|
||
//! projected out and its level fixed on exit (anchor or mean zero) on its
|
||
//! own, so enclosed fluid pockets or a body that splits the domain cannot
|
||
//! make the iteration diverge on a globally-compatible source.
|
||
//! - Limit: the coarsening is isotropic point-aggregation, so the
|
||
//! grid-independence holds for cells of aspect ratio near one. Measured
|
||
//! (review): aspect 2 → 6 iterations at every size, aspect 4 → 15–22,
|
||
//! aspect 10 → 37–91 (32² → 256²), aspect 100 → 405 at 256². Strongly
|
||
//! anisotropic grids need semi-coarsening or line smoothing; the flow
|
||
//! solvers here use square cells.
|
||
//! - The smoother is symmetric Gauss–Seidel (forward then backward sweep),
|
||
//! so each level's smoother is self-adjoint in the `A` inner product and
|
||
//! the V-cycle with `R = Pᵀ` is a symmetric preconditioner, which keeps
|
||
//! CG valid. The coarsest level is solved by 50 symmetric GS sweeps.
|
||
//! - On a pure-Neumann problem the mean over active cells is projected out
|
||
//! of the right-hand side and of every preconditioned residual, so CG runs
|
||
//! on the range of `A` and the level of `p` is fixed on exit (anchor cell
|
||
//! or mean zero).
|
||
//! - An active cell whose row is empty (`ap == 0`: an isolated cell with no
|
||
//! correctable face and no Dirichlet contribution) has no equation: the
|
||
//! solver never reads or writes it and it is excluded from the residual
|
||
//! — the same rule the embedded SOR applied.
|
||
//!
|
||
//! The convergence stop is the L1 norm of the TRUE residual `b - A p` over
|
||
//! the active cells: the CG recurrence residual is checked every iteration
|
||
//! and, when it passes the tolerance, the true residual is recomputed and
|
||
//! must pass too before the solve reports convergence.
|
||
|
||
/// Cell-centred five-point symmetric positive (semi-)definite problem
|
||
///
|
||
/// ```text
|
||
/// ap_i p_i - ae_i p_E - aw_i p_W - an_i p_N - as_i p_S = rhs_i on active cells,
|
||
/// ```
|
||
///
|
||
/// with `ap_i = ae_i + aw_i + an_i + as_i + extra_diag_i`. Row-major index
|
||
/// `j * nx + i`. The coefficient to a neighbour is 0 when that face is
|
||
/// prescribed or the neighbour is inactive. Symmetry is required:
|
||
/// `ae[j*nx+i] == aw[j*nx+i+1]`, `an[j*nx+i] == as_[(j+1)*nx+i]`.
|
||
/// `extra_diag` holds Dirichlet (outlet) contributions that have no
|
||
/// neighbour.
|
||
#[derive(Debug, Clone)]
|
||
pub struct PoissonProblem {
|
||
/// Cells in x.
|
||
pub nx: usize,
|
||
/// Cells in y.
|
||
pub ny: usize,
|
||
/// Active (fluid) cells; inactive cells carry no equation and are never
|
||
/// read or written by the solver.
|
||
pub active: Vec<bool>,
|
||
/// Coefficient to the east neighbour `(j, i+1)`.
|
||
pub ae: Vec<f64>,
|
||
/// Coefficient to the west neighbour `(j, i-1)`.
|
||
pub aw: Vec<f64>,
|
||
/// Coefficient to the north neighbour `(j+1, i)`.
|
||
pub an: Vec<f64>,
|
||
/// Coefficient to the south neighbour `(j-1, i)`.
|
||
pub as_: Vec<f64>,
|
||
/// Diagonal-only (Dirichlet) contribution.
|
||
pub extra_diag: Vec<f64>,
|
||
/// Right-hand side.
|
||
pub rhs: Vec<f64>,
|
||
}
|
||
|
||
impl PoissonProblem {
|
||
/// All cells active, every coefficient and the right-hand side zero.
|
||
#[must_use]
|
||
pub fn new(nx: usize, ny: usize) -> Self {
|
||
let n = nx * ny;
|
||
Self {
|
||
nx,
|
||
ny,
|
||
active: vec![true; n],
|
||
ae: vec![0.0; n],
|
||
aw: vec![0.0; n],
|
||
an: vec![0.0; n],
|
||
as_: vec![0.0; n],
|
||
extra_diag: vec![0.0; n],
|
||
rhs: vec![0.0; n],
|
||
}
|
||
}
|
||
|
||
/// Row-major index of cell `(j, i)`.
|
||
#[inline]
|
||
#[must_use]
|
||
pub fn index(&self, j: usize, i: usize) -> usize {
|
||
j * self.nx + i
|
||
}
|
||
|
||
/// Diagonal coefficient `ap` of cell `idx`.
|
||
#[inline]
|
||
#[must_use]
|
||
pub fn diagonal(&self, idx: usize) -> f64 {
|
||
self.ae[idx] + self.aw[idx] + self.an[idx] + self.as_[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)
|
||
}
|
||
|
||
/// Sum over the active cells that carry an equation (`ap > 0`) of
|
||
/// `|rhs - (ap p - Σ a_nb p_nb)|`. Neighbour values are read only from
|
||
/// active in-range neighbours.
|
||
#[must_use]
|
||
pub fn residual_l1(&self, p: &[f64]) -> f64 {
|
||
let (nx, ny) = (self.nx, self.ny);
|
||
let mut sum = 0.0;
|
||
for j in 0..ny {
|
||
for i in 0..nx {
|
||
let idx = j * nx + 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];
|
||
}
|
||
sum += (self.rhs[idx] - (ap * p[idx] - nb)).abs();
|
||
}
|
||
}
|
||
sum
|
||
}
|
||
|
||
/// Checks array lengths, non-negativity of every coefficient, zero
|
||
/// coefficients across domain edges and towards inactive cells, and
|
||
/// symmetry to `1e-12` relative.
|
||
pub fn validate(&self) -> Result<(), String> {
|
||
let n = self.nx * self.ny;
|
||
for (name, len) in [
|
||
("active", self.active.len()),
|
||
("ae", self.ae.len()),
|
||
("aw", self.aw.len()),
|
||
("an", self.an.len()),
|
||
("as_", self.as_.len()),
|
||
("extra_diag", self.extra_diag.len()),
|
||
("rhs", self.rhs.len()),
|
||
] {
|
||
if len != n {
|
||
return Err(format!("{name}: length {len}, expected nx*ny = {n}"));
|
||
}
|
||
}
|
||
let (nx, ny) = (self.nx, self.ny);
|
||
for j in 0..ny {
|
||
for i in 0..nx {
|
||
let idx = j * nx + i;
|
||
for (name, v) in [
|
||
("ae", self.ae[idx]),
|
||
("aw", self.aw[idx]),
|
||
("an", self.an[idx]),
|
||
("as_", self.as_[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 symmetric = |a: f64, b: f64| (a - b).abs() <= 1e-12 * a.abs().max(b.abs());
|
||
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]
|
||
));
|
||
}
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
/// Arithmetic precision of the multigrid PRECONDITIONER (the conjugate
|
||
/// gradient around it, its true residual, the mean projections and the
|
||
/// stopping rule stay `f64`). `F32` is the registered mixed-precision
|
||
/// design of `overset_metal_campaign.md` §3.2 M1 / §5.2: it decides
|
||
/// whether an fp32-only device (Apple Metal) can carry the pressure solve
|
||
/// of a coupled march. Default `F64` = bit-identical to the historical
|
||
/// solver.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||
pub enum MgPrecision {
|
||
/// Double precision throughout (the default; bit-identical).
|
||
#[default]
|
||
F64,
|
||
/// Single-precision V-cycle inside the double-precision CG.
|
||
F32,
|
||
}
|
||
|
||
/// The V-cycle's smoother ordering (PERF-2, `docs/perf2_campaign.md`).
|
||
/// `Lexicographic` is the recorded regime (row-major symmetric
|
||
/// Gauss–Seidel, a dependency chain through the division per cell);
|
||
/// `RedBlack` updates the two colours of the five-point stencil in turn —
|
||
/// each colour a map of independent cells (threads, vectors, the GPU) —
|
||
/// and is a different preconditioner, gated by the noise probe and the
|
||
/// anchor's band, never bit-identical.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||
pub enum MgSmoother {
|
||
#[default]
|
||
Lexicographic,
|
||
RedBlack,
|
||
}
|
||
|
||
/// Multigrid preconditioner parameters.
|
||
#[derive(Debug, Clone)]
|
||
pub struct MultigridParameters {
|
||
/// Precision of the V-cycle (see [`MgPrecision`]).
|
||
pub precision: MgPrecision,
|
||
/// Smoother ordering (see [`MgSmoother`]).
|
||
pub smoother: MgSmoother,
|
||
/// Threads for the red-black colour maps, the residual and the matvec
|
||
/// (rayon; default 1 = the serial code path). The parallel path computes
|
||
/// the same per-cell values from the same inputs and sums in the same
|
||
/// order, so it is bit-identical to the serial one. No effect on the
|
||
/// lexicographic smoother (a dependency chain).
|
||
pub threads: usize,
|
||
/// Symmetric Gauss–Seidel sweeps before AND after the coarse correction
|
||
/// (default 2; a value of 0 is treated as 1). One count for both on
|
||
/// purpose: unequal pre/post counts make the V-cycle non-symmetric and
|
||
/// conjugate gradient invalid — measured symmetry defect 0.9 for (1, 0)
|
||
/// against 1e-14 for (1, 1) — so the API does not let it happen.
|
||
pub smoother_sweeps: usize,
|
||
/// Stop coarsening when the active cells are at most this many; that
|
||
/// level is solved by 50 symmetric GS sweeps (default 32).
|
||
pub coarsest_cells: usize,
|
||
/// CG iteration cap (default 500).
|
||
pub max_iterations: usize,
|
||
}
|
||
|
||
impl Default for MultigridParameters {
|
||
fn default() -> Self {
|
||
Self {
|
||
precision: MgPrecision::F64,
|
||
smoother: MgSmoother::Lexicographic,
|
||
threads: 1,
|
||
smoother_sweeps: 2,
|
||
coarsest_cells: 32,
|
||
max_iterations: 500,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Outcome of a [`solve_multigrid_pcg`] call. `converged == false` means
|
||
/// the returned `p` did not reach the tolerance (iteration cap, rounding
|
||
/// floor, or an inconsistent system) — callers must act on it, which is
|
||
/// why the type is `#[must_use]`.
|
||
#[derive(Debug, Clone, Copy)]
|
||
#[must_use]
|
||
pub struct PoissonSolution {
|
||
/// CG iterations performed.
|
||
pub iterations: usize,
|
||
/// L1 true residual `Σ |b - A p|` over the active cells at exit.
|
||
pub residual: f64,
|
||
/// `residual < tolerance` at exit.
|
||
pub converged: bool,
|
||
/// Wall time of the setup (hierarchy, fine level, components) [ns].
|
||
pub setup_ns: u64,
|
||
/// Wall time of the CG iteration (including the V-cycles) [ns].
|
||
pub iterate_ns: u64,
|
||
}
|
||
|
||
/// Which inner solver a projection uses for its pressure-correction system.
|
||
///
|
||
/// `Sor` is the historical point successive over-relaxation at the optimal
|
||
/// Poisson factor; `Multigrid` is [`solve_multigrid_pcg`] on the same
|
||
/// coefficients, the same right-hand side and the same true-residual stop.
|
||
/// The two land on the same discrete pressure correction (to the inner
|
||
/// tolerance); they differ only in cost, which is mesh-independent for
|
||
/// multigrid and grows with the mesh for SOR.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||
pub enum PoissonSolverKind {
|
||
/// Point SOR, `omega = 2 / (1 + sin(pi / max(nx, ny)))`, 2000-sweep cap.
|
||
#[default]
|
||
Sor,
|
||
/// Conjugate gradient preconditioned by one geometric multigrid V-cycle.
|
||
Multigrid,
|
||
}
|
||
|
||
/// Symmetric GS sweeps on the coarsest level.
|
||
const COARSEST_SWEEPS: usize = 50;
|
||
/// Over-correction factor applied to the prolongated coarse correction.
|
||
///
|
||
/// With piecewise-constant prolongation and summation restriction the
|
||
/// Galerkin coarse operator of an aggregate-by-2 hierarchy is exactly twice
|
||
/// the geometric `2h` operator (a coarse face collects `2^(d-1)` fine faces
|
||
/// whose conductance is `h^(d-2)` each, against `(2h)^(d-2)` for the `2h`
|
||
/// cell), while the restricted residual of a smooth error is the full
|
||
/// `(2h)^d f`: the plain Galerkin correction is therefore half the geometric
|
||
/// one for the smooth modes, in any dimension and for any uniform face
|
||
/// conductance — the well-known under-correction of unsmoothed aggregation
|
||
/// (Braess 1995). Scaling the correction by 2 restores the geometric
|
||
/// correction for the smooth modes. The preconditioner stays symmetric
|
||
/// (a scalar factor), and stays positive definite: the coarse correction
|
||
/// operator `ω P B_c R A` is `A`-self-adjoint and positive, so the V-cycle
|
||
/// error propagator `S (I − ω P B_c R A) S` has `A`-spectrum below 1 and
|
||
/// `M⁻¹A = I − E` is positive on the range of `A`. Measured: the PCG
|
||
/// iteration count for a 1e-8 residual reduction on the Neumann box goes
|
||
/// from `[14, 20, 29, 40]` (32²…256², growing a level at a time) to
|
||
/// `[4, 4, 4, 4]`, and on the ragged mask from `[20, 28, 39]` to `[8, 8, 8]`.
|
||
const COARSE_CORRECTION: f64 = 2.0;
|
||
/// Hard cap on the number of levels (a 1×1 coarsest is reached long before).
|
||
const MAX_LEVELS: usize = 64;
|
||
|
||
/// The scalar the V-cycle runs in. `f64` reproduces the historical solver
|
||
/// bit for bit (same operations in the same order on the same values);
|
||
/// `f32` is the mixed-precision probe.
|
||
pub trait MgScalar:
|
||
Copy
|
||
+ Send
|
||
+ Sync
|
||
+ PartialEq
|
||
+ PartialOrd
|
||
+ std::ops::Add<Output = Self>
|
||
+ std::ops::Sub<Output = Self>
|
||
+ std::ops::Mul<Output = Self>
|
||
+ std::ops::Div<Output = Self>
|
||
+ std::ops::AddAssign
|
||
+ Send
|
||
+ Sync
|
||
+ 'static
|
||
{
|
||
/// Zero.
|
||
const ZERO: Self;
|
||
/// From double.
|
||
fn from_f64(v: f64) -> Self;
|
||
/// To double.
|
||
fn to_f64(self) -> f64;
|
||
/// Absolute value.
|
||
fn abs(self) -> Self;
|
||
}
|
||
|
||
impl MgScalar for f64 {
|
||
const ZERO: Self = 0.0;
|
||
#[inline]
|
||
fn from_f64(v: f64) -> Self {
|
||
v
|
||
}
|
||
#[inline]
|
||
fn to_f64(self) -> f64 {
|
||
self
|
||
}
|
||
#[inline]
|
||
fn abs(self) -> Self {
|
||
f64::abs(self)
|
||
}
|
||
}
|
||
|
||
impl MgScalar for f32 {
|
||
const ZERO: Self = 0.0;
|
||
#[inline]
|
||
fn from_f64(v: f64) -> Self {
|
||
v as f32
|
||
}
|
||
#[inline]
|
||
fn to_f64(self) -> f64 {
|
||
f64::from(self)
|
||
}
|
||
#[inline]
|
||
fn abs(self) -> Self {
|
||
f32::abs(self)
|
||
}
|
||
}
|
||
|
||
/// One level of the hierarchy: the problem (kept in `f64` for coarsening
|
||
/// and inspection), its coefficients and diagonal in the V-cycle scalar,
|
||
/// the list of active cells that carry an equation (row-major), and the
|
||
/// parent map into the next coarser level.
|
||
struct Level<T: MgScalar> {
|
||
problem: PoissonProblem,
|
||
/// `problem.active && ap > 0`: cells with an equation.
|
||
active: Vec<bool>,
|
||
ae: Vec<T>,
|
||
aw: Vec<T>,
|
||
an: Vec<T>,
|
||
as_: Vec<T>,
|
||
ap: Vec<T>,
|
||
/// Row-major indices of the active cells.
|
||
cells: Vec<usize>,
|
||
/// The active cells with `i + j` even / odd, each row-major (the
|
||
/// red-black smoother's two independent maps on the five-point stencil).
|
||
red: Vec<usize>,
|
||
black: Vec<usize>,
|
||
/// Fine index → coarse index (empty on the coarsest level).
|
||
coarse_of: Vec<usize>,
|
||
}
|
||
|
||
/// V-cycle work vectors of one level.
|
||
struct Work<T: MgScalar> {
|
||
/// Right-hand side of the residual equation on this level.
|
||
b: Vec<T>,
|
||
/// Correction on this level.
|
||
x: Vec<T>,
|
||
/// Residual.
|
||
r: Vec<T>,
|
||
/// Scratch for the parallel maps (one entry per active cell).
|
||
tmp: Vec<T>,
|
||
}
|
||
|
||
impl<T: MgScalar> Work<T> {
|
||
fn new(n: usize) -> Self {
|
||
Self {
|
||
b: vec![T::ZERO; n],
|
||
x: vec![T::ZERO; n],
|
||
r: vec![T::ZERO; n],
|
||
tmp: vec![T::ZERO; n],
|
||
}
|
||
}
|
||
}
|
||
|
||
impl<T: MgScalar> Level<T> {
|
||
fn new(mut problem: PoissonProblem) -> Self {
|
||
let (nx, ny) = (problem.nx, problem.ny);
|
||
let n = nx * ny;
|
||
let ap: Vec<f64> = (0..n).map(|idx| problem.diagonal(idx)).collect();
|
||
let active: Vec<bool> = (0..n)
|
||
.map(|idx| problem.active[idx] && ap[idx] > 0.0)
|
||
.collect();
|
||
// Sanitise: a coefficient towards an inactive or out-of-range
|
||
// neighbour is never used, so zero it — the stencil can then branch
|
||
// on the coefficient alone and never touches an inactive value.
|
||
for j in 0..ny {
|
||
for i in 0..nx {
|
||
let idx = j * nx + i;
|
||
if !active[idx] {
|
||
problem.ae[idx] = 0.0;
|
||
problem.aw[idx] = 0.0;
|
||
problem.an[idx] = 0.0;
|
||
problem.as_[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;
|
||
}
|
||
}
|
||
}
|
||
let cells: Vec<usize> = (0..n).filter(|&idx| active[idx]).collect();
|
||
let parity = |idx: usize| (idx % nx + idx / nx) % 2;
|
||
let red: Vec<usize> = cells
|
||
.iter()
|
||
.copied()
|
||
.filter(|&idx| parity(idx) == 0)
|
||
.collect();
|
||
let black: Vec<usize> = cells
|
||
.iter()
|
||
.copied()
|
||
.filter(|&idx| parity(idx) == 1)
|
||
.collect();
|
||
let cast = |v: &[f64]| v.iter().map(|&x| T::from_f64(x)).collect::<Vec<T>>();
|
||
Self {
|
||
ae: cast(&problem.ae),
|
||
aw: cast(&problem.aw),
|
||
an: cast(&problem.an),
|
||
as_: cast(&problem.as_),
|
||
ap: cast(&ap),
|
||
problem,
|
||
active,
|
||
cells,
|
||
red,
|
||
black,
|
||
coarse_of: Vec::new(),
|
||
}
|
||
}
|
||
|
||
/// `Σ a_nb x_nb` for cell `idx`; only active in-range neighbours are
|
||
/// read (their coefficients are the only non-zero ones after
|
||
/// sanitising).
|
||
#[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];
|
||
}
|
||
s
|
||
}
|
||
|
||
/// `y = A x` on the active cells.
|
||
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);
|
||
}
|
||
}
|
||
|
||
/// `r = b - A x` on the active cells; returns its L1 norm.
|
||
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
|
||
}
|
||
|
||
/// [`Self::residual`] with the per-cell values computed in parallel into
|
||
/// `tmp` (one per active cell, in `cells` order) and the L1 sum taken
|
||
/// serially in the same order: bit-identical.
|
||
fn residual_par(&self, b: &[T], x: &[T], r: &mut [T], tmp: &mut [T]) -> T {
|
||
use rayon::prelude::*;
|
||
let n = self.cells.len();
|
||
self.cells
|
||
.par_iter()
|
||
.zip(tmp[..n].par_iter_mut())
|
||
.with_min_len(PAR_MIN_LEN)
|
||
.for_each(|(&idx, t)| {
|
||
*t = b[idx] - (self.ap[idx] * x[idx] - self.neighbour_sum(x, idx));
|
||
});
|
||
let mut l1 = T::ZERO;
|
||
for (k, &idx) in self.cells.iter().enumerate() {
|
||
let v = tmp[k];
|
||
r[idx] = v;
|
||
l1 += v.abs();
|
||
}
|
||
l1
|
||
}
|
||
|
||
/// One red-black half-sweep over `colour` in parallel: every cell of a
|
||
/// colour reads only the other colour, so the values are computed into
|
||
/// `tmp` from the unchanged `x` and written back — the serial
|
||
/// half-sweep's values exactly.
|
||
fn half_sweep_par(&self, colour: &[usize], b: &[T], x: &mut [T], tmp: &mut [T]) {
|
||
use rayon::prelude::*;
|
||
let n = colour.len();
|
||
{
|
||
let x_ro: &[T] = x;
|
||
colour
|
||
.par_iter()
|
||
.zip(tmp[..n].par_iter_mut())
|
||
.with_min_len(PAR_MIN_LEN)
|
||
.for_each(|(&idx, t)| {
|
||
*t = (b[idx] + self.neighbour_sum(x_ro, idx)) / self.ap[idx];
|
||
});
|
||
}
|
||
for (k, &idx) in colour.iter().enumerate() {
|
||
x[idx] = tmp[k];
|
||
}
|
||
}
|
||
|
||
/// [`Self::symmetric_gs_rb`] on threads (bit-identical).
|
||
fn symmetric_gs_rb_par(&self, b: &[T], x: &mut [T], tmp: &mut [T]) {
|
||
self.half_sweep_par(&self.red, b, x, tmp);
|
||
self.half_sweep_par(&self.black, b, x, tmp);
|
||
self.half_sweep_par(&self.black, b, x, tmp);
|
||
self.half_sweep_par(&self.red, b, x, tmp);
|
||
}
|
||
|
||
/// One symmetric Gauss–Seidel sweep (forward then backward) on `A x = b`.
|
||
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];
|
||
}
|
||
}
|
||
|
||
/// One symmetric red-black Gauss–Seidel sweep: red, black, black, red —
|
||
/// within a colour every cell reads only the other colour, so each
|
||
/// half-sweep is a map (symmetric as a preconditioner, like the
|
||
/// lexicographic pair).
|
||
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];
|
||
}
|
||
}
|
||
|
||
/// One symmetric sweep in the chosen ordering (`tmp`: the parallel
|
||
/// path's scratch, used only with red-black on threads).
|
||
fn smooth(&self, b: &[T], x: &mut [T], tmp: &mut [T], smoother: MgSmoother, parallel: bool) {
|
||
match (smoother, parallel) {
|
||
(MgSmoother::Lexicographic, _) => self.symmetric_gs(b, x),
|
||
(MgSmoother::RedBlack, false) => self.symmetric_gs_rb(b, x),
|
||
(MgSmoother::RedBlack, true) => self.symmetric_gs_rb_par(b, x, tmp),
|
||
}
|
||
}
|
||
|
||
/// Galerkin coarsening: coarse face coefficient = sum of the fine
|
||
/// coefficients across that coarse face, coarse `extra_diag` = sum of
|
||
/// the children's. Returns the coarse problem and the parent map.
|
||
fn coarsen(&self) -> (PoissonProblem, Vec<usize>) {
|
||
let (nx, ny) = (self.problem.nx, self.problem.ny);
|
||
let nxc = (nx / 2).max(1);
|
||
let nyc = (ny / 2).max(1);
|
||
let cx = |i: usize| (i / 2).min(nxc - 1);
|
||
let cy = |j: usize| (j / 2).min(nyc - 1);
|
||
let mut coarse = PoissonProblem::new(nxc, nyc);
|
||
coarse.active.fill(false);
|
||
let mut coarse_of = vec![usize::MAX; nx * ny];
|
||
for &idx in &self.cells {
|
||
let (i, j) = (idx % nx, idx / nx);
|
||
let (ic, jc) = (cx(i), cy(j));
|
||
let c = jc * nxc + ic;
|
||
coarse_of[idx] = c;
|
||
coarse.active[c] = true;
|
||
coarse.extra_diag[c] += self.problem.extra_diag[idx];
|
||
// Coefficients towards inactive neighbours are already zero.
|
||
if self.problem.ae[idx] != 0.0 && cx(i + 1) != ic {
|
||
coarse.ae[c] += self.problem.ae[idx];
|
||
}
|
||
if self.problem.aw[idx] != 0.0 && cx(i - 1) != ic {
|
||
coarse.aw[c] += self.problem.aw[idx];
|
||
}
|
||
if self.problem.an[idx] != 0.0 && cy(j + 1) != jc {
|
||
coarse.an[c] += self.problem.an[idx];
|
||
}
|
||
if self.problem.as_[idx] != 0.0 && cy(j - 1) != jc {
|
||
coarse.as_[c] += self.problem.as_[idx];
|
||
}
|
||
}
|
||
(coarse, coarse_of)
|
||
}
|
||
}
|
||
|
||
/// The multigrid hierarchy: level 0 is the fine problem.
|
||
pub(crate) struct Hierarchy<T: MgScalar = f64> {
|
||
levels: Vec<Level<T>>,
|
||
work: Vec<Work<T>>,
|
||
sweeps: usize,
|
||
smoother: MgSmoother,
|
||
parallel: bool,
|
||
}
|
||
|
||
impl<T: MgScalar> Hierarchy<T> {
|
||
/// Builds the hierarchy down to `coarsest_cells` active cells (or until
|
||
/// coarsening stops reducing the count).
|
||
pub(crate) fn build(problem: &PoissonProblem, params: &MultigridParameters) -> Self {
|
||
let mut levels = vec![Level::<T>::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 = Level::<T>::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| Work::new(l.problem.nx * l.problem.ny))
|
||
.collect();
|
||
Self {
|
||
levels,
|
||
work,
|
||
sweeps: params.smoother_sweeps.max(1),
|
||
smoother: params.smoother,
|
||
parallel: params.threads > 1,
|
||
}
|
||
}
|
||
|
||
/// Number of levels.
|
||
pub(crate) fn depth(&self) -> usize {
|
||
self.levels.len()
|
||
}
|
||
|
||
/// The problem on level `l` (level 0 is the fine problem, with the
|
||
/// coefficients towards inactive cells sanitised to zero).
|
||
pub(crate) fn problem(&self, l: usize) -> &PoissonProblem {
|
||
&self.levels[l].problem
|
||
}
|
||
|
||
/// Parent map from level `l` to level `l + 1` (`usize::MAX` on cells
|
||
/// without an equation).
|
||
pub(crate) fn coarse_of(&self, l: usize) -> &[usize] {
|
||
&self.levels[l].coarse_of
|
||
}
|
||
|
||
/// Row-major indices of the active cells of level `l`.
|
||
pub(crate) fn cells(&self, l: usize) -> &[usize] {
|
||
&self.levels[l].cells
|
||
}
|
||
|
||
/// `z = M⁻¹ r`: one V-cycle on `A z = r` from `z = 0`. Only active
|
||
/// entries of `r` are read and only active entries of `z` are written.
|
||
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]);
|
||
}
|
||
// Down: smooth from zero, restrict the residual.
|
||
let (smoother, parallel) = (self.smoother, self.parallel);
|
||
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 Work { b, x, r, tmp } = wf;
|
||
let par_here = parallel && fine.cells.len() >= PAR_MIN_CELLS;
|
||
for _ in 0..self.sweeps {
|
||
fine.smooth(b, x, tmp, smoother, par_here);
|
||
}
|
||
if par_here {
|
||
fine.residual_par(b, x, r, tmp);
|
||
} else {
|
||
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];
|
||
// A coarse cell without an equation (a whole component
|
||
// inside one aggregate) receives the component's zero sum.
|
||
if coarse.active[c] {
|
||
wc.b[c] += r[idx];
|
||
}
|
||
}
|
||
}
|
||
// Coarsest: a fixed number of symmetric sweeps from zero.
|
||
{
|
||
let bottom = &levels[depth - 1];
|
||
let wb = &mut self.work[depth - 1];
|
||
for &idx in &bottom.cells {
|
||
wb.x[idx] = T::ZERO;
|
||
}
|
||
let Work { b, x, tmp, .. } = wb;
|
||
// The coarsest level is tiny (≤ 32 cells): always serial.
|
||
for _ in 0..COARSEST_SWEEPS {
|
||
bottom.smooth(b, x, tmp, smoother, false);
|
||
}
|
||
}
|
||
// Up: prolongate, smooth.
|
||
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 Work { b, x, tmp, .. } = wf;
|
||
let par_here = parallel && fine.cells.len() >= PAR_MIN_CELLS;
|
||
for _ in 0..self.sweeps {
|
||
fine.smooth(b, x, tmp, smoother, par_here);
|
||
}
|
||
}
|
||
for &idx in &levels[0].cells {
|
||
z[idx] = self.work[0].x[idx].to_f64();
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Conjugate gradient on the active cells preconditioned by one V-cycle of
|
||
/// geometric multigrid (see the module documentation for the hierarchy).
|
||
///
|
||
/// Stops when the L1 true residual `Σ |b - A p|` over the active cells is
|
||
/// below `tolerance` (absolute — the caller passes a scale-relative value)
|
||
/// or after `params.max_iterations` iterations.
|
||
///
|
||
/// Singular (pure Neumann) problems: CG runs on the consistent system (the
|
||
/// caller guarantees a compatible rhs to rounding; the mean over active
|
||
/// cells is projected out of the rhs to be safe, and the residual reported
|
||
/// and tested is against that projected rhs — an incompatible rhs shows up
|
||
/// as the gap between it and `problem.residual_l1(p)`), and on exit `p` is shifted
|
||
/// so `p[anchor] == 0` when `anchor` is `Some` and names an active cell,
|
||
/// else to mean zero over the active cells. On a non-singular problem the
|
||
/// level is determined by the equations and `anchor` is ignored.
|
||
///
|
||
/// `p` is the initial guess and the result; inactive entries are neither
|
||
/// read nor written.
|
||
pub fn solve_multigrid_pcg(
|
||
problem: &PoissonProblem,
|
||
p: &mut [f64],
|
||
params: &MultigridParameters,
|
||
tolerance: f64,
|
||
anchor: Option<usize>,
|
||
) -> PoissonSolution {
|
||
match params.precision {
|
||
MgPrecision::F64 => solve_pcg_with::<f64>(problem, p, params, tolerance, anchor),
|
||
MgPrecision::F32 => solve_pcg_with::<f32>(problem, p, params, tolerance, anchor),
|
||
}
|
||
}
|
||
|
||
/// The CG driver: `f64` throughout, with the V-cycle in `T`. The fine
|
||
/// level used for the CG's own products and true residual is a separate
|
||
/// `f64` level, so the precision of the preconditioner never enters the
|
||
/// stopping rule or the reported residual.
|
||
/// The operator part of a [`PoissonProblem`] (everything but the
|
||
/// right-hand side) plus the hierarchy parameters, kept to decide whether
|
||
/// a prepared solver can be reused (PERF-2 P1.1, `docs/perf2_campaign.md`).
|
||
/// The comparison is exact (bit patterns), so a reuse changes nothing.
|
||
struct OperatorKey {
|
||
nx: usize,
|
||
ny: usize,
|
||
active: Vec<bool>,
|
||
coefficients: Vec<u64>,
|
||
smoother_sweeps: usize,
|
||
coarsest_cells: usize,
|
||
smoother: MgSmoother,
|
||
threads: usize,
|
||
}
|
||
|
||
impl OperatorKey {
|
||
fn of(problem: &PoissonProblem, params: &MultigridParameters) -> Self {
|
||
let coefficients = problem
|
||
.ae
|
||
.iter()
|
||
.chain(&problem.aw)
|
||
.chain(&problem.an)
|
||
.chain(&problem.as_)
|
||
.chain(&problem.extra_diag)
|
||
.map(|v| v.to_bits())
|
||
.collect();
|
||
Self {
|
||
nx: problem.nx,
|
||
ny: problem.ny,
|
||
active: problem.active.clone(),
|
||
coefficients,
|
||
smoother_sweeps: params.smoother_sweeps,
|
||
coarsest_cells: params.coarsest_cells,
|
||
smoother: params.smoother,
|
||
threads: params.threads,
|
||
}
|
||
}
|
||
|
||
fn matches(&self, problem: &PoissonProblem, params: &MultigridParameters) -> bool {
|
||
self.nx == problem.nx
|
||
&& self.ny == problem.ny
|
||
&& self.smoother_sweeps == params.smoother_sweeps
|
||
&& self.coarsest_cells == params.coarsest_cells
|
||
&& self.smoother == params.smoother
|
||
&& self.threads == params.threads
|
||
&& self.active == problem.active
|
||
&& self.coefficients.iter().copied().eq(problem
|
||
.ae
|
||
.iter()
|
||
.chain(&problem.aw)
|
||
.chain(&problem.an)
|
||
.chain(&problem.as_)
|
||
.chain(&problem.extra_diag)
|
||
.map(|v| v.to_bits()))
|
||
}
|
||
}
|
||
|
||
/// Everything the CG driver derives from the OPERATOR: the hierarchy, the
|
||
/// `f64` fine level, the active cells and the connected components. A
|
||
/// deterministic function of the operator; the work vectors inside the
|
||
/// hierarchy are re-initialised on the active set at every use, so a
|
||
/// prepared solver reused for another right-hand side gives the same
|
||
/// answer as a fresh one, bit for bit.
|
||
struct Prepared<T: MgScalar> {
|
||
key: OperatorKey,
|
||
hier: Hierarchy<T>,
|
||
fine: Level<f64>,
|
||
cells: Vec<usize>,
|
||
components: Components,
|
||
}
|
||
|
||
impl<T: MgScalar> Prepared<T> {
|
||
fn build(problem: &PoissonProblem, params: &MultigridParameters) -> Self {
|
||
let hier = Hierarchy::<T>::build(problem, params);
|
||
let fine = Level::<f64>::new(problem.clone());
|
||
let cells: Vec<usize> = fine.cells.clone();
|
||
let components = Components::find(problem, &cells);
|
||
Self {
|
||
key: OperatorKey::of(problem, params),
|
||
hier,
|
||
fine,
|
||
cells,
|
||
components,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Threads pay only when a level is large: below this many active cells a
|
||
/// level's maps run serially even with `threads > 1` (the fork-join of a
|
||
/// map over a few thousand cells costs more than the map), and each task
|
||
/// covers at least [`PAR_MIN_LEN`] cells. Both are pure scheduling: the
|
||
/// values are the serial ones.
|
||
const PAR_MIN_CELLS: usize = 8192;
|
||
const PAR_MIN_LEN: usize = 2048;
|
||
|
||
/// Configure rayon's global thread pool for the multigrid's parallel maps
|
||
/// (PERF-2 P2). Idempotent: a pool that already exists is kept (rayon
|
||
/// refuses a second global pool) — the first caller decides.
|
||
pub fn configure_threads(threads: usize) -> usize {
|
||
let n = threads.max(1);
|
||
if n > 1 {
|
||
let _ = rayon::ThreadPoolBuilder::new()
|
||
.num_threads(n)
|
||
.build_global();
|
||
}
|
||
n
|
||
}
|
||
|
||
/// PERF-2 P3 (`docs/perf2_campaign.md`): one hierarchy level exported for a
|
||
/// device V-cycle — the sanitised f32 coefficients, the active-cell and
|
||
/// colour index lists, the parent map, and the coarse cells' children in
|
||
/// CSR form (so a restriction can sum in a fixed order on the device).
|
||
#[derive(Debug, Clone)]
|
||
pub struct LevelExport {
|
||
pub nx: usize,
|
||
pub ny: usize,
|
||
pub cells: Vec<u32>,
|
||
pub red: Vec<u32>,
|
||
pub black: Vec<u32>,
|
||
/// Fine cell → coarse cell (`u32::MAX` without an equation; empty on
|
||
/// the coarsest level).
|
||
pub coarse_of: Vec<u32>,
|
||
/// For the NEXT level's cells, in its `cells` order: the fine cells
|
||
/// restricting into each (CSR: `children_ptr[c]..children_ptr[c + 1]`).
|
||
pub children_ptr: Vec<u32>,
|
||
pub children_idx: Vec<u32>,
|
||
pub ae: Vec<f32>,
|
||
pub aw: Vec<f32>,
|
||
pub an: Vec<f32>,
|
||
pub as_: Vec<f32>,
|
||
pub ap: Vec<f32>,
|
||
}
|
||
|
||
/// The f32 hierarchy of `problem` (red-black colour lists included), level
|
||
/// 0 fine, for a device implementation of [`Hierarchy::apply_preconditioner`].
|
||
pub fn export_hierarchy(problem: &PoissonProblem, params: &MultigridParameters) -> Vec<LevelExport> {
|
||
let hier = Hierarchy::<f32>::build(problem, params);
|
||
let depth = hier.levels.len();
|
||
(0..depth)
|
||
.map(|l| {
|
||
let lv = &hier.levels[l];
|
||
let to_u32 = |v: &[usize]| v.iter().map(|&i| i as u32).collect::<Vec<u32>>();
|
||
let (children_ptr, children_idx) = if l + 1 < depth {
|
||
let coarse = &hier.levels[l + 1];
|
||
let mut pos = vec![usize::MAX; coarse.problem.nx * coarse.problem.ny];
|
||
for (k, &c) in coarse.cells.iter().enumerate() {
|
||
pos[c] = k;
|
||
}
|
||
let mut lists: Vec<Vec<u32>> = vec![Vec::new(); coarse.cells.len()];
|
||
for &idx in &lv.cells {
|
||
let c = lv.coarse_of[idx];
|
||
if coarse.active[c] {
|
||
lists[pos[c]].push(idx as u32);
|
||
}
|
||
}
|
||
let mut ptr = Vec::with_capacity(lists.len() + 1);
|
||
let mut flat = Vec::new();
|
||
ptr.push(0u32);
|
||
for list in &lists {
|
||
flat.extend_from_slice(list);
|
||
ptr.push(flat.len() as u32);
|
||
}
|
||
(ptr, flat)
|
||
} else {
|
||
(Vec::new(), Vec::new())
|
||
};
|
||
LevelExport {
|
||
nx: lv.problem.nx,
|
||
ny: lv.problem.ny,
|
||
cells: to_u32(&lv.cells),
|
||
red: to_u32(&lv.red),
|
||
black: to_u32(&lv.black),
|
||
coarse_of: lv
|
||
.coarse_of
|
||
.iter()
|
||
.map(|&c| if c == usize::MAX { u32::MAX } else { c as u32 })
|
||
.collect(),
|
||
children_ptr,
|
||
children_idx,
|
||
ae: lv.ae.clone(),
|
||
aw: lv.aw.clone(),
|
||
an: lv.an.clone(),
|
||
as_: lv.as_.clone(),
|
||
ap: lv.ap.clone(),
|
||
}
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// One f32 V-cycle of `problem`'s hierarchy on `r` (the CPU reference for a
|
||
/// device V-cycle): `z = M⁻¹ r` exactly as the preconditioner computes it.
|
||
pub fn vcycle_f32_reference(
|
||
problem: &PoissonProblem,
|
||
params: &MultigridParameters,
|
||
r: &[f64],
|
||
z: &mut [f64],
|
||
) {
|
||
let mut hier = Hierarchy::<f32>::build(problem, params);
|
||
hier.apply_preconditioner(r, z);
|
||
}
|
||
|
||
/// A reusable prepared solver per V-cycle precision (PERF-2 P1.1): the
|
||
/// operator's hierarchy is rebuilt only when the operator changes.
|
||
#[derive(Default)]
|
||
pub struct PcgCache {
|
||
f64: Option<Prepared<f64>>,
|
||
f32: Option<Prepared<f32>>,
|
||
}
|
||
|
||
impl PcgCache {
|
||
/// Number of prepared operators held (0, 1 or 2).
|
||
pub fn len(&self) -> usize {
|
||
usize::from(self.f64.is_some()) + usize::from(self.f32.is_some())
|
||
}
|
||
|
||
/// Whether nothing is cached yet.
|
||
pub fn is_empty(&self) -> bool {
|
||
self.len() == 0
|
||
}
|
||
}
|
||
|
||
/// [`solve_multigrid_pcg`] with the operator's hierarchy taken from
|
||
/// `cache` when the operator (coefficients, active mask, hierarchy
|
||
/// parameters) is bit-identical to the cached one, rebuilt into it
|
||
/// otherwise. The answer is that of the uncached solve, bit for bit.
|
||
pub fn solve_multigrid_pcg_cached(
|
||
problem: &PoissonProblem,
|
||
p: &mut [f64],
|
||
params: &MultigridParameters,
|
||
tolerance: f64,
|
||
anchor: Option<usize>,
|
||
cache: &mut PcgCache,
|
||
) -> PoissonSolution {
|
||
match params.precision {
|
||
MgPrecision::F64 => {
|
||
solve_cached_with::<f64>(problem, p, params, tolerance, anchor, &mut cache.f64)
|
||
}
|
||
MgPrecision::F32 => {
|
||
solve_cached_with::<f32>(problem, p, params, tolerance, anchor, &mut cache.f32)
|
||
}
|
||
}
|
||
}
|
||
|
||
fn solve_cached_with<T: MgScalar>(
|
||
problem: &PoissonProblem,
|
||
p: &mut [f64],
|
||
params: &MultigridParameters,
|
||
tolerance: f64,
|
||
anchor: Option<usize>,
|
||
slot: &mut Option<Prepared<T>>,
|
||
) -> 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(Prepared::<T>::build(problem, params));
|
||
}
|
||
let setup_ns = t_entry.elapsed().as_nanos() as u64;
|
||
let prep = slot.as_mut().expect("prepared");
|
||
run_pcg(prep, problem, p, params, tolerance, anchor, setup_ns)
|
||
}
|
||
|
||
fn solve_pcg_with<T: MgScalar>(
|
||
problem: &PoissonProblem,
|
||
p: &mut [f64],
|
||
params: &MultigridParameters,
|
||
tolerance: f64,
|
||
anchor: Option<usize>,
|
||
) -> PoissonSolution {
|
||
let t_entry = std::time::Instant::now();
|
||
let mut prep = Prepared::<T>::build(problem, params);
|
||
let setup_ns = t_entry.elapsed().as_nanos() as u64;
|
||
run_pcg(&mut prep, problem, p, params, tolerance, anchor, setup_ns)
|
||
}
|
||
|
||
/// The CG loop on a prepared operator (see [`Prepared`]).
|
||
fn run_pcg<T: MgScalar>(
|
||
prep: &mut Prepared<T>,
|
||
problem: &PoissonProblem,
|
||
p: &mut [f64],
|
||
params: &MultigridParameters,
|
||
tolerance: f64,
|
||
anchor: Option<usize>,
|
||
setup_ns: u64,
|
||
) -> PoissonSolution {
|
||
let n = problem.nx * problem.ny;
|
||
assert_eq!(p.len(), n, "p must have nx*ny entries");
|
||
debug_assert!(
|
||
problem.validate().is_ok(),
|
||
"invalid PoissonProblem: {:?}",
|
||
problem.validate()
|
||
);
|
||
let hier = &mut prep.hier;
|
||
let fine = &prep.fine;
|
||
let cells: &[usize] = &prep.cells;
|
||
let components = &prep.components;
|
||
let active_n = cells.len();
|
||
if active_n == 0 {
|
||
return PoissonSolution {
|
||
iterations: 0,
|
||
residual: 0.0,
|
||
converged: true,
|
||
setup_ns,
|
||
iterate_ns: 0,
|
||
};
|
||
}
|
||
|
||
// Connected components of the active cells (through faces with a
|
||
// non-zero coefficient). A component with no Dirichlet contribution is
|
||
// a pure-Neumann block of the system and singular ON ITS OWN: its
|
||
// right-hand side must have zero mean for the block to be consistent,
|
||
// whatever the other components carry. A single global mean projection
|
||
// is not enough — two enclosed pockets with opposite imbalances sum to
|
||
// zero globally and still make CG diverge, and the divergence pollutes
|
||
// even a well-posed Dirichlet component (found in review, 1e-8
|
||
// relative was enough). So the mean is projected per singular
|
||
// component, and the exit shift is applied per singular component.
|
||
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::<f64>() / 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::<f64>();
|
||
let l1 = |a: &[f64]| cells.iter().map(|&idx| a[idx].abs()).sum::<f64>();
|
||
|
||
// Right-hand side (per-component mean projected out where singular).
|
||
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: &Level<f64>| -> 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| {
|
||
// Level of each singular component: the anchor's component is
|
||
// shifted so p[anchor] == 0, every other singular component to mean
|
||
// zero. Non-singular components have their level fixed by their
|
||
// Dirichlet data and are left alone.
|
||
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::<f64>() / 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);
|
||
}
|
||
|
||
hier.apply_preconditioner(&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;
|
||
// True residual at the last resynchronisation: when a resynchronised
|
||
// true residual no longer improves on the previous one, the recurrence
|
||
// has hit the rounding floor of `b - A p` and further iterations cannot
|
||
// reach the tolerance — stop, honestly unconverged.
|
||
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 {
|
||
// Breakdown (r = 0 to rounding, or a non-positive curvature
|
||
// from rounding in the null space): stop on the true residual.
|
||
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 {
|
||
// The recurrence residual passed: confirm against the TRUE
|
||
// residual, and resynchronise if rounding has let them drift.
|
||
res = true_residual(p, &mut r, &fine);
|
||
if res < tolerance || res > 0.9 * last_true {
|
||
return finish(p, iterations, res);
|
||
}
|
||
last_true = res;
|
||
}
|
||
hier.apply_preconditioner(&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)
|
||
}
|
||
|
||
/// Connected components of the active cells of a [`PoissonProblem`],
|
||
/// connected through faces carrying a non-zero coefficient, and whether
|
||
/// each component is singular (carries no Dirichlet contribution).
|
||
struct Components {
|
||
/// Component id per cell (`usize::MAX` for inactive cells).
|
||
id: Vec<usize>,
|
||
/// Member cells per component.
|
||
members: Vec<Vec<usize>>,
|
||
/// Per component: no active member has `extra_diag > 0`.
|
||
singular: Vec<bool>,
|
||
}
|
||
|
||
impl Components {
|
||
fn find(problem: &PoissonProblem, cells: &[usize]) -> Self {
|
||
let (nx, ny) = (problem.nx, problem.ny);
|
||
let n = nx * ny;
|
||
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 (j, i) = (idx / 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]);
|
||
}
|
||
}
|
||
members.push(list);
|
||
singular.push(!has_dirichlet);
|
||
}
|
||
Self {
|
||
id,
|
||
members,
|
||
singular,
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests;
|