From 82bb448d89881c24eee33559a7f506e466ed11e8 Mon Sep 17 00:00:00 2001 From: Kristof Schroeder Date: Wed, 23 Sep 2026 18:59:20 +0200 Subject: [PATCH 1/2] fix(lsmr): normalize each vector by its own norm, never through 1/x MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A subnormal α, β or ‖b‖ has no normal reciprocal, and ‖A‖ = f64::MAX leaves 1/α₁ subnormal, so `v *= 1/α` failed or zeroed the basis vector. `normalize` divides when the reciprocal leaves the range (`drscl`), at every stream site: init/restart, both steps, the MGK ring and its p̃ update, which divides p̃ by α first when β/α itself leaves the range (`dlascl`). --- CHANGELOG.md | 2 +- crates/schwarz-precond/src/lsmr/bidiag.rs | 59 ++++++------- .../schwarz-precond/src/lsmr/bidiag/tests.rs | 13 ++- crates/schwarz-precond/src/lsmr/tests.rs | 1 + .../schwarz-precond/src/lsmr/tests/range.rs | 86 +++++++++++++++++++ 5 files changed, 123 insertions(+), 38 deletions(-) create mode 100644 crates/schwarz-precond/src/lsmr/tests/range.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index a7270eb5..470bceeb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,7 +42,7 @@ and this project follows [Semantic Versioning](https://semver.org/). - A `design` that is neither a 2-D `uint32` array nor a list of `Effect` raised `ValueError` where the documented type is `TypeError`, and `AdditiveSchwarz` accepted a wrong-type `local_solver` at construction, deferring the `TypeError` to solve time (#248). - LSMR no longer certifies a stop whose recurrence estimates collapsed. Tolerance stops are checked against `‖b − A x‖`, a non-finite `α`, `β`, `⟨v, Mv⟩`, or `‖b‖` fails with `SolveError::InvalidInput`, and an overflowing or subnormal `‖A‖` no longer zeroes the normal-equation ratio; a failed check reports `LsmrStopReason::FalseConvergence` with `converged = false` and the returned iterate's normal-equation residual (#290, #297, #303, #362). - A warm-started solve measures its residuals against the original `b`, including at a budget stop and where `‖Aᵀb‖²` overflows. -- An `α` or `β` whose square underflows is recovered from a scaled norm, where LSMR reported `x = 0` converged. +- An `α` or `β` whose square underflows is recovered from a scaled norm, where LSMR reported `x = 0` converged; a subnormal `α`, `β` or `‖b‖` no longer fails with `SolveError::InvalidInput`. - LSMR's solution update no longer drops out when `‖A‖` is small (`≈ 1e-9`) or beyond `1e±154`, where it reported `x = 0` converged, and a solve whose `‖A‖²` or `‖A‖‖b‖` leaves the double range no longer fails with `SolveError::InvalidInput`, misses its normal-equation stop, or reports a zero normal-equation residual. - A preconditioned solve with local reorthogonalization (`local_size`) could fail with `SolveError::InvalidInput` ("preconditioner not positive definite") on a positive definite preconditioner, once reorthogonalization cancelled the last Krylov direction to rounding noise; the pair is now recomputed before the check. - LSMR's residual estimate is its own `‖r_k‖` rather than LSQR's smaller `|φ̄_k|`, which let `ResidualTolerance` fire before the tolerance was met. diff --git a/crates/schwarz-precond/src/lsmr/bidiag.rs b/crates/schwarz-precond/src/lsmr/bidiag.rs index ede27397..28977217 100644 --- a/crates/schwarz-precond/src/lsmr/bidiag.rs +++ b/crates/schwarz-precond/src/lsmr/bidiag.rs @@ -62,12 +62,15 @@ pub(super) fn axpby(y: &mut [f64], x: &[f64], alpha: f64, beta: f64) { } } -/// In-place scalar multiply `y *= s`. Parallel above the threshold. +/// `y /= d` without over/underflow unless `y / d` does, to one rounding at `f64::MAX` (`drscl`). #[inline] -fn scale_in_place(y: &mut [f64], s: f64) { +fn normalize(y: &mut [f64], d: f64) { + let inv = 1.0 / d; let seq = |c: &mut [f64]| { - for yi in c { - *yi *= s; + if inv.is_normal() { + c.iter_mut().for_each(|yi| *yi *= inv); + } else { + c.iter_mut().for_each(|yi| *yi /= d); } }; if y.len() >= LSMR_PAR_THRESHOLD { @@ -175,7 +178,7 @@ pub(super) fn residual_into( /// Scales `u` to unit length given its already-computed `β = ‖u‖`; a zero `u` stays zero. fn scale_to_unit(u: &mut [f64], beta: f64) { if beta > 0.0 { - scale_in_place(u, 1.0 / beta); + normalize(u, beta); } } @@ -277,16 +280,16 @@ impl WindowRing<2> { } } - /// Copy normalized `v` and `p_tilde · inv_alpha` into the next slots, advancing the ring. - fn push(&mut self, v: &[f64], p_tilde_unscaled: &[f64], inv_alpha: f64) { + /// Copy normalized `v` and `p_tilde / alpha` into the next slots, advancing the ring. + fn push(&mut self, v: &[f64], p_tilde_unscaled: &[f64], alpha: f64) { let slot = self.advance(); self.lane_mut(0, slot).copy_from_slice(v); - for (dst, &src) in self - .lane_mut(1, slot) - .iter_mut() - .zip(p_tilde_unscaled.iter()) - { - *dst = src * inv_alpha; + let p = self.lane_mut(1, slot); + if alpha > 0.0 { + p.copy_from_slice(p_tilde_unscaled); + normalize(p, alpha); + } else { + p.fill(0.0); } } } @@ -313,7 +316,7 @@ pub(super) fn metric_gradient_norm( } // `M⁻¹` is linear; the raw gradient's metric product overflows past `‖Aᵀr‖ ≈ 1e154`. let scale = if plain.is_normal() { plain } else { 1.0 }; - scale_in_place(g, 1.0 / scale); + normalize(g, scale); preconditioner.apply(g, mv)?; Ok(scale * alpha_from_vp(mv, g)?) } @@ -351,8 +354,7 @@ impl Bidiagonalization for GolubKahan<'_, A> { self.alpha = 0.0; return Ok(BidiagStep { alpha: 0.0, beta }); } - // beta > 0 here: the beta == 0 lucky breakdown returned above. - scale_in_place(&mut self.bufs.u, 1.0 / beta); + normalize(&mut self.bufs.u, beta); self.operator .apply_adjoint(&self.bufs.u, &mut self.bufs.atu)?; @@ -365,7 +367,7 @@ impl Bidiagonalization for GolubKahan<'_, A> { } let alpha = finite(alpha, "α")?; if alpha > 0.0 { - scale_in_place(&mut self.bufs.v, 1.0 / alpha); + normalize(&mut self.bufs.v, alpha); } if let Some(reorth) = &mut self.bufs.local_reorth { @@ -394,7 +396,7 @@ impl Bidiagonalization for GolubKahan<'_, A> { .apply_adjoint(&self.bufs.u, &mut self.bufs.v)?; let alpha = finite(par_norm(&self.bufs.v), "α")?; if alpha > 0.0 { - scale_in_place(&mut self.bufs.v, 1.0 / alpha); + normalize(&mut self.bufs.v, alpha); } if let Some(reorth) = &mut self.bufs.local_reorth { reorth.clear(); @@ -460,12 +462,11 @@ impl Bidiagonalization .apply(&self.bufs.p_tilde, &mut self.bufs.v)?; let alpha = alpha_from_vp(&self.bufs.v, &self.bufs.p_tilde)?; if alpha > 0.0 { - scale_in_place(&mut self.bufs.v, 1.0 / alpha); + normalize(&mut self.bufs.v, alpha); } if let Some(reorth) = &mut self.bufs.local_reorth { reorth.clear(); - let inv_alpha = if alpha > 0.0 { 1.0 / alpha } else { 0.0 }; - reorth.push(&self.bufs.v, &self.bufs.p_tilde, inv_alpha); + reorth.push(&self.bufs.v, &self.bufs.p_tilde, alpha); } self.alpha = alpha; self.beta_prev_inv = 1.0; // u was normalized @@ -609,7 +610,12 @@ impl<'a, A: Operator + ?Sized, M: Operator + ?Sized> ModifiedGolubKahan<'a, A, M self.alpha > 0.0, "self.alpha must be > 0; lsmr_from_bidiag's loop guard prevents step() after alpha=0", ); - let p_coeff = beta / self.alpha; + let mut p_coeff = beta / self.alpha; + if !p_coeff.is_normal() { + // `β/α` can leave the range while `p̃ / α = M v` does not, so divide first (`dlascl`). + normalize(&mut self.bufs.p_tilde, self.alpha); + p_coeff = beta; + } axpby(&mut self.bufs.p_tilde, &self.bufs.atu, beta_inv, -p_coeff); Ok(()) } @@ -634,16 +640,11 @@ impl<'a, A: Operator + ?Sized, M: Operator + ?Sized> ModifiedGolubKahan<'a, A, M }; if alpha_new > 0.0 { - scale_in_place(&mut self.bufs.v, 1.0 / alpha_new); + normalize(&mut self.bufs.v, alpha_new); } if let Some(reorth) = &mut self.bufs.local_reorth { - let inv_alpha = if alpha_new > 0.0 { - 1.0 / alpha_new - } else { - 0.0 - }; - reorth.push(&self.bufs.v, &self.bufs.p_tilde, inv_alpha); + reorth.push(&self.bufs.v, &self.bufs.p_tilde, alpha_new); } Ok(alpha_new) diff --git a/crates/schwarz-precond/src/lsmr/bidiag/tests.rs b/crates/schwarz-precond/src/lsmr/bidiag/tests.rs index eac4973d..0c26f0fb 100644 --- a/crates/schwarz-precond/src/lsmr/bidiag/tests.rs +++ b/crates/schwarz-precond/src/lsmr/bidiag/tests.rs @@ -31,14 +31,11 @@ fn alpha_from_vp_clamps_a_pair_within_root_epsilon(#[case] v: &[f64], #[case] p: assert_eq!(alpha_from_vp(v, p).expect("within √ε"), 0.0, "{v:?}·{p:?}"); } -/// `beta == 0.0` and `alpha > 0.0` are both false for NaN; unguarded, an overflow poisons the run. -#[rstest] -fn a_non_finite_operator_norm_is_an_error(#[values(f64::NAN, f64::MAX)] bad: f64) { - let result = crate::lsmr::lsmr(&DiagOp(vec![bad, 1.0]), &[1.0, 1.0], 1e-10, 50, None); - assert!( - matches!(result, Err(SolveError::InvalidInput { .. })), - "{bad:e} accepted" - ); +/// `beta == 0.0` and `alpha > 0.0` are both false for NaN; unguarded, it poisons the run. +#[test] +fn a_non_finite_operator_norm_is_an_error() { + let result = crate::lsmr::lsmr(&DiagOp(vec![f64::NAN, 1.0]), &[1.0, 1.0], 1e-10, 50, None); + assert!(matches!(result, Err(SolveError::InvalidInput { .. }))); } /// A finite adjoint keeps `init` clean, so the overflow reaches the `step` guard on β. diff --git a/crates/schwarz-precond/src/lsmr/tests.rs b/crates/schwarz-precond/src/lsmr/tests.rs index fa0a7864..45edcf44 100644 --- a/crates/schwarz-precond/src/lsmr/tests.rs +++ b/crates/schwarz-precond/src/lsmr/tests.rs @@ -3,4 +3,5 @@ mod audit; mod breakdown; mod ladder; +mod range; mod solve; diff --git a/crates/schwarz-precond/src/lsmr/tests/range.rs b/crates/schwarz-precond/src/lsmr/tests/range.rs new file mode 100644 index 00000000..143397e6 --- /dev/null +++ b/crates/schwarz-precond/src/lsmr/tests/range.rs @@ -0,0 +1,86 @@ +//! Extreme scales: no vector leaves the double range unless the answer itself does. + +use rstest::rstest; + +use super::super::*; +use crate::lsmr::fixtures::*; +use crate::Operator; + +#[derive(Clone, Copy, Debug)] +enum Metric { + None, + Identity, + Diagonal, +} + +fn solve( + a: &A, + b: &[f64], + metric: Metric, + local_size: Option, + warm_start: Option<&[f64]>, +) -> Result { + let n = a.ncols(); + let diagonal = DiagOp( + (0..n) + .map(|j| if j % 2 == 0 { 4.0 } else { 0.25 }) + .collect(), + ); + let options = MlsmrOptions { + warm_start, + local_size, + ..Default::default() + }; + match metric { + Metric::None => lsmr(a, b, 1e-10, 50, local_size), + Metric::Identity => mlsmr(a, b, &IdentityOp { n }, 1e-10, 50, options), + Metric::Diagonal => mlsmr(a, b, &diagonal, 1e-10, 50, options), + } +} + +/// `α₁ = ‖Aᵀb‖ / ‖b‖ ≈ 1e-310` is subnormal although `‖A‖` and `‖b‖` are not, so `1/α₁` is ∞. +#[rstest] +fn a_subnormal_initial_gradient_normalizes_its_basis_vector( + #[values(Metric::None, Metric::Identity, Metric::Diagonal)] metric: Metric, + #[values(None, Some(2))] local_size: Option, +) { + let r = solve( + &DiagOp(vec![0.0, 1.0]), + &[1e200, 1e-110], + metric, + local_size, + None, + ) + .expect("subnormal-gradient solve"); + + assert!(r.converged, "{:?}", r.stop_reason); + assert_eq!(r.x[0], 0.0); + assert!((r.x[1] / 1e-110 - 1.0).abs() < 1e-9, "{:?}", r.x); +} + +/// A subnormal `‖b‖` has no representable reciprocal either. +#[rstest] +fn a_subnormal_rhs_normalizes_its_first_vector( + #[values(Metric::None, Metric::Identity, Metric::Diagonal)] metric: Metric, +) { + let r = solve(&IdentityOp { n: 2 }, &[1e-310, 0.0], metric, None, None) + .expect("subnormal-rhs solve"); + + assert!(r.converged, "{:?}", r.stop_reason); + assert!((r.x[0] / 1e-310 - 1.0).abs() < 1e-9, "{:?}", r.x); + assert_eq!(r.x[1], 0.0); +} + +/// `‖A‖ = f64::MAX` leaves `1/α₁` subnormal; a stop dropping column 2 must hold its backward error. +#[test] +fn an_operator_at_the_top_of_the_range_solves_to_its_backward_error() { + let a = DiagOp(vec![f64::MAX, 1.0]); + let b = [1.0, 1.0]; + let r = lsmr(&a, &b, 1e-10, 50, None).expect("top-of-range solve"); + + assert!(r.converged, "{:?}", r.stop_reason); + assert!((r.x[0] * f64::MAX - 1.0).abs() < 1e-9, "{:?}", r.x); + let normr = vec_norm(&[b[0] - f64::MAX * r.x[0], b[1] - r.x[1]]); + let backward_error = normal_equation_residual(&a, &r.x, &b) / f64::MAX / normr; + assert!(backward_error <= 1e-10, "{backward_error:e}"); +} From c1f634e32c2389762b06dfc05fea95d474ab53c8 Mon Sep 17 00:00:00 2001 From: Kristof Schroeder Date: Thu, 24 Sep 2026 14:51:34 +0200 Subject: [PATCH 2/2] refactor(lsmr): let normalize own the zero-norm guard, drop scale_to_unit --- crates/schwarz-precond/src/lsmr/bidiag.rs | 49 +++++++++-------------- 1 file changed, 18 insertions(+), 31 deletions(-) diff --git a/crates/schwarz-precond/src/lsmr/bidiag.rs b/crates/schwarz-precond/src/lsmr/bidiag.rs index 28977217..d374174c 100644 --- a/crates/schwarz-precond/src/lsmr/bidiag.rs +++ b/crates/schwarz-precond/src/lsmr/bidiag.rs @@ -65,18 +65,20 @@ pub(super) fn axpby(y: &mut [f64], x: &[f64], alpha: f64, beta: f64) { /// `y /= d` without over/underflow unless `y / d` does, to one rounding at `f64::MAX` (`drscl`). #[inline] fn normalize(y: &mut [f64], d: f64) { - let inv = 1.0 / d; - let seq = |c: &mut [f64]| { - if inv.is_normal() { - c.iter_mut().for_each(|yi| *yi *= inv); + if d > 0.0 { + let inv = 1.0 / d; + let seq = |c: &mut [f64]| { + if inv.is_normal() { + c.iter_mut().for_each(|yi| *yi *= inv); + } else { + c.iter_mut().for_each(|yi| *yi /= d); + } + }; + if y.len() >= LSMR_PAR_THRESHOLD { + y.par_chunks_mut(LSMR_UPDATE_CHUNK).for_each(seq); } else { - c.iter_mut().for_each(|yi| *yi /= d); + seq(y); } - }; - if y.len() >= LSMR_PAR_THRESHOLD { - y.par_chunks_mut(LSMR_UPDATE_CHUNK).for_each(seq); - } else { - seq(y); } } @@ -175,13 +177,6 @@ pub(super) fn residual_into( Ok(axpy_with_norm(u, rhs, -1.0)) } -/// Scales `u` to unit length given its already-computed `β = ‖u‖`; a zero `u` stays zero. -fn scale_to_unit(u: &mut [f64], beta: f64) { - if beta > 0.0 { - normalize(u, beta); - } -} - /// Ring of recent basis vectors for windowed MGS; the disabled state is `None`, so `cap > 0`. struct WindowRing { /// `L` flat buffers; lane `l`, slot `s` is `[s*n .. s*n + n]` of `lanes[l]`. @@ -366,9 +361,7 @@ impl Bidiagonalization for GolubKahan<'_, A> { alpha = par_norm(&self.bufs.v); } let alpha = finite(alpha, "α")?; - if alpha > 0.0 { - normalize(&mut self.bufs.v, alpha); - } + normalize(&mut self.bufs.v, alpha); if let Some(reorth) = &mut self.bufs.local_reorth { reorth.push(&self.bufs.v); @@ -391,13 +384,11 @@ impl Bidiagonalization for GolubKahan<'_, A> { } fn restart(&mut self, beta: f64) -> Result { - scale_to_unit(&mut self.bufs.u, beta); + normalize(&mut self.bufs.u, beta); self.operator .apply_adjoint(&self.bufs.u, &mut self.bufs.v)?; let alpha = finite(par_norm(&self.bufs.v), "α")?; - if alpha > 0.0 { - normalize(&mut self.bufs.v, alpha); - } + normalize(&mut self.bufs.v, alpha); if let Some(reorth) = &mut self.bufs.local_reorth { reorth.clear(); reorth.push(&self.bufs.v); @@ -455,15 +446,13 @@ impl Bidiagonalization } fn restart(&mut self, beta: f64) -> Result { - scale_to_unit(&mut self.bufs.u, beta); + normalize(&mut self.bufs.u, beta); self.operator .apply_adjoint(&self.bufs.u, &mut self.bufs.p_tilde)?; self.preconditioner .apply(&self.bufs.p_tilde, &mut self.bufs.v)?; let alpha = alpha_from_vp(&self.bufs.v, &self.bufs.p_tilde)?; - if alpha > 0.0 { - normalize(&mut self.bufs.v, alpha); - } + normalize(&mut self.bufs.v, alpha); if let Some(reorth) = &mut self.bufs.local_reorth { reorth.clear(); reorth.push(&self.bufs.v, &self.bufs.p_tilde, alpha); @@ -639,9 +628,7 @@ impl<'a, A: Operator + ?Sized, M: Operator + ?Sized> ModifiedGolubKahan<'a, A, M alpha => alpha?, }; - if alpha_new > 0.0 { - normalize(&mut self.bufs.v, alpha_new); - } + normalize(&mut self.bufs.v, alpha_new); if let Some(reorth) = &mut self.bufs.local_reorth { reorth.push(&self.bufs.v, &self.bufs.p_tilde, alpha_new);