fix(lsmr): normalize each vector by its own norm, never through 1/x

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.

crates/schwarz-precond/src/lsmr/bidiag.rs

@@ -62,18 +62,23 @@ 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) {
-    let seq = |c: &mut [f64]| {
-        for yi in c {
-            *yi *= s;
+fn normalize(y: &mut [f64], d: f64) {
+    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 {
+            seq(y);
         }
-    };
-    if y.len() >= LSMR_PAR_THRESHOLD {
-        y.par_chunks_mut(LSMR_UPDATE_CHUNK).for_each(seq);
-    } else {
-        seq(y);
     }
 }
 
@@ -172,13 +177,6 @@ pub(super) fn residual_into<A: Operator + ?Sized>(
     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 {
-        scale_in_place(u, 1.0 / beta);
-    }
-}
-
 /// Ring of recent basis vectors for windowed MGS; the disabled state is `None`, so `cap > 0`.
 struct WindowRing<const L: usize> {
     /// `L` flat buffers; lane `l`, slot `s` is `[s*n .. s*n + n]` of `lanes[l]`.
@@ -277,16 +275,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 +311,7 @@ pub(super) fn metric_gradient_norm<A: Operator + ?Sized, M: Operator + ?Sized>(
     }
     // `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 +349,7 @@ impl<A: Operator + ?Sized> 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)?;
@@ -364,9 +361,7 @@ impl<A: Operator + ?Sized> Bidiagonalization for GolubKahan<'_, A> {
             alpha = par_norm(&self.bufs.v);
         }
         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 {
             reorth.push(&self.bufs.v);
@@ -389,13 +384,11 @@ impl<A: Operator + ?Sized> Bidiagonalization for GolubKahan<'_, A> {
     }
 
     fn restart(&mut self, beta: f64) -> Result<BidiagStep, SolveError> {
-        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 {
-            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();
             reorth.push(&self.bufs.v);
@@ -453,19 +446,16 @@ impl<A: Operator + ?Sized, M: Operator + ?Sized> Bidiagonalization
     }
 
     fn restart(&mut self, beta: f64) -> Result<BidiagStep, SolveError> {
-        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 {
-            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 +599,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(())
     }
@@ -633,17 +628,10 @@ impl<'a, A: Operator + ?Sized, M: Operator + ?Sized> ModifiedGolubKahan<'a, A, M
             alpha => alpha?,
         };
 
-        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)

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 β.

crates/schwarz-precond/src/lsmr/tests.rs

@@ -3,4 +3,5 @@
 mod audit;
 mod breakdown;
 mod ladder;
+mod range;
 mod solve;

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: Operator>(
+    a: &A,
+    b: &[f64],
+    metric: Metric,
+    local_size: Option<usize>,
+    warm_start: Option<&[f64]>,
+) -> Result<LsmrResult, SolveError> {
+    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<usize>,
+) {
+    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}");
+}