fix(lsmr): run the recurrence on u₁ and test its stop in exact exponents

CHANGELOG.md

@@ -43,7 +43,7 @@ and this project follows [Semantic Versioning](https://semver.org/).
 - 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.
-- LSMR's solution update no longer drops out when `‖A‖` is small (`≈ 1e-9`) or beyond `1e±154`, where it reported `x = 0` converged.
+- 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.
 - A non-finite or negative `ScalingConfig::tolerance` silently disabled the dominance certificate under both failure policies, since every comparison against it is `>`; it is now rejected as `BuildError::InvalidScalingTolerance`.

crates/schwarz-precond/src/lsmr.rs

@@ -7,6 +7,7 @@
 mod bidiag;
 #[cfg(test)]
 mod fixtures;
+mod magnitude;
 mod recurrence;
 #[cfg(test)]
 mod tests;
@@ -18,6 +19,7 @@ use bidiag::{
     axpby, metric_gradient_norm, residual_into, BidiagStep, Bidiagonalization, GolubKahan,
     ModifiedGolubKahan,
 };
+use magnitude::Magnitude;
 use recurrence::{ConvergenceCriteria, LsmrRecurrenceState, RotationStep, SolutionState, Stop};
 
 /// Euclidean norm of a vector.
@@ -333,7 +335,7 @@ pub fn mlsmr<A: Operator + ?Sized, M: Operator + ?Sized>(
         ModifiedGolubKahan::init(operator, preconditioner, &rhs, rhs_norm, local_size)?;
     let warm_start = warm_start.zip(metric).map(|(x0, metric)| WarmStart {
         x0,
-        reference: NormalEqReference::warm(metric, step1),
+        reference: NormalEqReference::warm(Magnitude::from(metric), step1),
     });
     let reference_norm = if b_norm > 0.0 { b_norm } else { rhs_norm };
     let criteria = ConvergenceCriteria::new(reference_norm, tol);
@@ -345,25 +347,26 @@ const MAX_RESTARTS: usize = 2;
 
 /// `‖Aᵀb‖` in the stream's metric, fixed for the solve so every pass reports against it.
 #[derive(Clone, Copy)]
-struct NormalEqReference(f64);
+struct NormalEqReference(Magnitude);
 
 impl NormalEqReference {
-    /// A cold stream's first `(α₁, β₁)` is `‖Aᵀb‖` already.
+    /// A cold stream's first `(α₁, β₁)` is `‖Aᵀb‖` already, though the product may not be a double.
     fn cold(step1: BidiagStep) -> Self {
-        Self((step1.alpha * step1.beta).abs().max(f64::MIN_POSITIVE))
+        Self(Magnitude::product(step1.alpha, step1.beta))
     }
 
-    /// A reference that carries no information leaves the stream's own `ζ̄₀` to divide.
-    fn warm(metric: f64, step1: BidiagStep) -> Self {
-        if metric > 0.0 && metric.is_finite() {
+    /// A reference that carries no information falls back to the cold `α₁β₁`.
+    fn warm(metric: Magnitude, step1: BidiagStep) -> Self {
+        if metric.is_normal() {
             Self(metric)
         } else {
             Self::cold(step1)
         }
     }
 
-    fn relative(self, estimate: f64) -> f64 {
-        estimate / self.0
+    fn relative(self, estimate: Magnitude) -> f64 {
+        debug_assert!(self.0.is_normal(), "a zero α₁ returns before any report");
+        (estimate / self.0).to_f64()
     }
 }
 
@@ -397,13 +400,17 @@ fn lsmr_from_bidiag<B: Bidiagonalization>(
             let (converged, normal_eq_residual) = match bidiag.hidden_gradient()? {
                 None => (true, 0.0),
                 Some(per_unit_residual) => {
-                    let normar = step1.beta * per_unit_residual;
+                    let normar = Magnitude::product(step1.beta, per_unit_residual);
                     let plain = bidiag.plain_gradient(b)?;
                     let a_norm_below = plain / vec_norm(b).max(f64::MIN_POSITIVE);
                     let informative = plain > 0.0 && plain.is_finite();
                     (
                         criteria.corroborates(step1.beta, normar, a_norm_below),
-                        if informative { normar / plain } else { 1.0 },
+                        if informative {
+                            (normar / Magnitude::from(plain)).to_f64()
+                        } else {
+                            1.0
+                        },
                     )
                 }
             };
@@ -426,7 +433,7 @@ fn lsmr_from_bidiag<B: Bidiagonalization>(
         let mut escalation = escalation.map(EscalationPolicy::handler);
         let mut convergence = criteria.start(step1.alpha);
         let mut recurrence = LsmrRecurrenceState::init(step1);
-        let mut solution = SolutionState::init(bidiag.v());
+        let mut solution = SolutionState::init(bidiag.v(), step1.beta);
         let mut prev_rot = RotationStep::initial();
         let stop_reason = 'pass: {
             while iterations < maxiter {
@@ -493,7 +500,7 @@ fn lsmr_from_bidiag<B: Bidiagonalization>(
                 // The estimate is the refuted claim; a seed from the staged residual measures `x`.
                 result.normal_eq_residual = if residual_norm.is_finite() {
                     let step = bidiag.restart(residual_norm)?;
-                    reference.relative(step.alpha * step.beta)
+                    reference.relative(Magnitude::product(step.alpha, step.beta))
                 } else {
                     residual_norm
                 };

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

@@ -0,0 +1,133 @@
+//! Norms as `m · 2^e`, for comparing gradients of problems whose scales share no range.
+
+use std::ops::Div;
+
+/// A nonnegative `m · 2^e`, `m ∈ [1, 2)`, zero or non-finite, for expressions a few factors deep.
+#[derive(Clone, Copy)]
+pub(super) struct Magnitude {
+    m: f64,
+    e: i32,
+}
+
+impl Magnitude {
+    /// `a · b`, exact in exponent where the `f64` product would over- or underflow.
+    pub(super) fn product(a: f64, b: f64) -> Self {
+        let (a, b) = (Self::from(a), Self::from(b));
+        Self::normalized(a.m * b.m, a.e + b.e)
+    }
+
+    /// The nearest `f64`, saturating to 0 or ∞; a subnormal result rounds a second time.
+    pub(super) fn to_f64(self) -> f64 {
+        if !self.m.is_normal() {
+            return self.m;
+        }
+        if self.e >= f64::MAX_EXP {
+            return f64::INFINITY;
+        }
+        if self.e >= f64::MIN_EXP - 1 {
+            return self.m * pow2(self.e);
+        }
+        // A subnormal result: the first factor is exact, so only the second rounds.
+        self.m * pow2((self.e + 64).max(f64::MIN_EXP - 1)) * pow2(-64)
+    }
+
+    /// Neither zero, infinite nor NaN.
+    pub(super) fn is_normal(self) -> bool {
+        self.m.is_normal()
+    }
+
+    fn normalized(m: f64, e: i32) -> Self {
+        if !m.is_normal() {
+            return Self { m, e: 0 };
+        }
+        let Self { m, e: shift } = Self::from(m);
+        Self { m, e: e + shift }
+    }
+}
+
+impl From<f64> for Magnitude {
+    fn from(x: f64) -> Self {
+        if !x.is_finite() || x == 0.0 {
+            return Self { m: x, e: 0 };
+        }
+        // A subnormal carries its leading bit below the mantissa field; 2^64 lifts it exactly.
+        let (x, bias) = if x.is_normal() {
+            (x, 0)
+        } else {
+            (x * pow2(64), -64)
+        };
+        debug_assert!(x > 0.0, "a norm, got {x:e}");
+        let bits = x.to_bits();
+        let e = ((bits >> 52) & 0x7ff) as i32 - 1023;
+        let m = f64::from_bits((bits & ((1 << 52) - 1)) | (1023 << 52));
+        Self { m, e: e + bias }
+    }
+}
+
+impl Div for Magnitude {
+    type Output = Self;
+    fn div(self, rhs: Self) -> Self {
+        Self::normalized(self.m / rhs.m, self.e - rhs.e)
+    }
+}
+
+/// `2^e` for a normal exponent, built from its bits since `powi`'s precision is unspecified.
+fn pow2(e: i32) -> f64 {
+    debug_assert!((f64::MIN_EXP - 1..f64::MAX_EXP).contains(&e));
+    f64::from_bits(((e + 1023) as u64) << 52)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Magnitude;
+    use rstest::rstest;
+
+    #[rstest]
+    #[case::one(1.0)]
+    #[case::large(3.7e300)]
+    #[case::smallest_normal(f64::MIN_POSITIVE)]
+    #[case::subnormal(1.3e-310)]
+    #[case::smallest_subnormal(f64::from_bits(1))]
+    #[case::largest(f64::MAX)]
+    #[case::zero(0.0)]
+    #[case::infinity(f64::INFINITY)]
+    fn a_float_round_trips(#[case] x: f64) {
+        assert_eq!(Magnitude::from(x).to_f64().to_bits(), x.to_bits());
+    }
+
+    /// Where the `f64` arithmetic stays in range the two agree bitwise.
+    #[rstest]
+    #[case::unit(3.0, 7.0, 5.0)]
+    #[case::large(3e150, 7e150, 5e-7)]
+    #[case::small(3e-150, 7e-150, 5e7)]
+    fn in_range_arithmetic_is_the_float_arithmetic(#[case] a: f64, #[case] b: f64, #[case] c: f64) {
+        let exact = Magnitude::product(a, b) / Magnitude::from(c);
+        assert_eq!(exact.to_f64().to_bits(), (a * b / c).to_bits());
+    }
+
+    /// A quotient back in range survives a product that is not.
+    #[rstest]
+    #[case::overflowing_product(1e300, 3e300, 1e300, 3e300)]
+    #[case::underflowing_product(1e-300, 3e-300, 1e-300, 3e-300)]
+    #[case::subnormal_quotient(1e-300, 3e-10, 1.0, 3e-310)]
+    fn out_of_range_products_divide_back(
+        #[case] a: f64,
+        #[case] b: f64,
+        #[case] c: f64,
+        #[case] expected: f64,
+    ) {
+        let got = (Magnitude::product(a, b) / Magnitude::from(c)).to_f64();
+        assert!(
+            (got / expected - 1.0).abs() < 1e-9,
+            "{got:e} vs {expected:e}"
+        );
+    }
+
+    #[test]
+    fn an_out_of_range_value_saturates() {
+        assert!(Magnitude::product(1e-300, 1e-300).is_normal());
+        assert_eq!(Magnitude::product(1e-300, 1e-300).to_f64(), 0.0);
+        assert_eq!(Magnitude::product(1e300, 1e300).to_f64(), f64::INFINITY);
+        assert!(Magnitude::from(f64::NAN).to_f64().is_nan());
+    }
+}