perf(lsmr): one true-residual check at a tolerance stop, restart from a refuted one

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project follows [Semantic Versioning](https://semver.org/).
 
 ### Changed
 
+- A tolerance stop is checked with one true-residual evaluation; a refuted one restarts from its iterate, at most twice, before `LsmrStopReason::FalseConvergence`.
 - Python `Effect` lists no longer deep-copy their level and slope buffers during design extraction; the binding borrows them while building off-GIL (#358).
 - Categorical `u32` labels need not be zero-based or contiguous: `Design` compacts observed labels to internal positions and `CoefficientLayout`/`CoefficientAddress` translate back, so gaps in sparse label ranges are neither allocated nor solved for (#228, #268).
 - **BREAKING:** Rust `Solver::new` takes `weights: Option<&[f64]>` instead of `Option<Vec<f64>>`, retaining only `W^{1/2}` in internal observation order. One-shot `solve`/`solve_batch` weights are unchanged.
@@ -18,7 +19,7 @@ and this project follows [Semantic Versioning](https://semver.org/).
 - **BREAKING:** `schwarz_precond::mlsmr` takes an `MlsmrOptions` in place of its trailing `local_size`.
 - **BREAKING:** `LsmrStopReason` gains `Escalated` and `WarmStartExact`, breaking exhaustive `match`es.
 - A warm start that already solves the system reports `WarmStartExact` instead of `ZeroRhs`.
-- The LSMR true-residual audit of a warm-started stop measures the total solution against the original `b`, anchoring its normal-equation leg to `‖Aᵀb‖` rather than the restart's initial residual; with `b = 0`, tolerances measure against `‖b − A x₀‖`.
+- A warm-started LSMR stop measures the total solution against the original `b`; with `b = 0`, tolerances measure against `‖b − A x₀‖`.
 - **BREAKING:** `ScalingConfig::max_sweeps` is now `max_iterations`, and `BuildWarning::UnscalableComponent` reports `iterations` in place of `sweeps`; the dominance certificate runs reduced CG, not relaxation sweeps.
 - **BREAKING:** The serialized `Preconditioner` wire format moved v12 → v17 (approx-chol 0.5, full construction config, build duration, `LocalSolverConfig::ridge`, and the built map's own Schwarz description in place of the strategy enum); 0.3.0 bytes no longer decode.
 - **BREAKING:** Coefficient addresses use caller-visible `u32` factor labels rather than internal `usize` level positions, affecting `CoefficientAddress::level` and the accepted range of Python coefficient-layout and unidentified-direction levels.
@@ -39,8 +40,8 @@ and this project follows [Semantic Versioning](https://semver.org/).
 
 - A design carrying varying slopes on two distinct factors could fail preconditioner construction with `matrix is not symmetric`, when rounding left the two triangles of the exact Schur complement unequal (#229).
 - 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 it has not solved. Tolerance stops are audited against the true residual and outside the preconditioner's metric, 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` (#290, #297, #303, #362).
-- A warm-started solve measures its residuals against the original `b`, including when it exhausts its iteration budget, and no longer fails a converged solve when the cold audit's metric product overflows at `‖Aᵀb‖ ≳ 1e154`.
+- 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 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.rs

@@ -15,7 +15,8 @@ use std::borrow::Cow;
 
 use crate::{Operator, SolveError};
 use bidiag::{
-    residual_into, BidiagStep, Bidiagonalization, Certificate, GolubKahan, ModifiedGolubKahan,
+    axpby, metric_gradient_norm, residual_into, BidiagStep, Bidiagonalization, GolubKahan,
+    ModifiedGolubKahan,
 };
 use recurrence::{ConvergenceCriteria, LsmrRecurrenceState, RotationStep, SolutionState, Stop};
 
@@ -42,13 +43,13 @@ pub(crate) fn vec_norm(v: &[f64]) -> f64 {
 pub struct LsmrResult {
     /// Solution vector.
     pub x: Vec<f64>,
-    /// Whether the solver converged within the tolerance.
+    /// Whether a tolerance stop matched `‖b − A x‖`, which misses drift within `range(A)`.
     pub converged: bool,
     /// Total number of iterations performed.
     pub iterations: usize,
-    /// Final residual norm estimate `‖b − A x‖`.
+    /// `‖b − A x‖`: recomputed at a tolerance stop, the recurrence's estimate at any other.
     pub residual_norm: f64,
-    /// Per-run relative normal-equation residual, measured in `M`'s metric when preconditioned.
+    /// `‖Âᵀ(b − A x)‖ / ‖Âᵀb‖`, measured in `M`'s metric when preconditioned.
     pub normal_eq_residual: f64,
     /// Reason the solver stopped.
     pub stop_reason: LsmrStopReason,
@@ -67,7 +68,7 @@ pub enum LsmrStopReason {
     NormalEquationTolerance,
     /// The warm start already solved the system: `b − A x0` was exactly zero.
     WarmStartExact,
-    /// A tolerance stop refuted by the true-residual audit; the recurrence estimates had collapsed.
+    /// A tolerance stop, and each restart from it, that `‖b − A x‖` refuted.
     FalseConvergence,
     /// The iteration budget was exhausted before convergence.
     MaxIterations,
@@ -244,7 +245,7 @@ pub fn lsmr<A: Operator + ?Sized>(
     lsmr_from_bidiag(bidiag, step1, b, None, criteria, maxiter, None)
 }
 
-/// Preconditioned LSMR with `M ≈ AᵀA` and one `M⁻¹` application per iteration.
+/// Preconditioned LSMR with `M ≈ AᵀA`, one `M⁻¹` apply per iteration; `M⁻¹` must be nonsingular.
 pub fn mlsmr<A: Operator + ?Sized, M: Operator + ?Sized>(
     operator: &A,
     b: &[f64],
@@ -313,145 +314,188 @@ pub fn mlsmr<A: Operator + ?Sized, M: Operator + ?Sized>(
         });
     }
 
+    // `‖Aᵀb‖` must be taken before the stream exists: the stream's own query clobbers `v₁`.
+    let metric = match warm_start {
+        None => None,
+        Some(_) => Some(metric_gradient_norm(
+            operator,
+            preconditioner,
+            b,
+            &mut vec![0.0; n],
+            &mut vec![0.0; n],
+        )?),
+    };
     let (bidiag, step1) =
         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),
+    });
     let reference_norm = if b_norm > 0.0 { b_norm } else { rhs_norm };
     let criteria = ConvergenceCriteria::new(reference_norm, tol);
-    lsmr_from_bidiag(
-        bidiag,
-        step1,
-        b,
-        warm_start,
-        criteria,
-        maxiter,
-        escalation.map(|policy| policy.handler()),
-    )
+    lsmr_from_bidiag(bidiag, step1, b, warm_start, criteria, maxiter, escalation)
 }
 
-/// A warm-started stream measures its references from the correction's own residual, so a cold
-/// certificate restores them in terms of `rhs`.
-fn audit_against_rhs<B: Bidiagonalization>(
-    bidiag: &mut B,
-    x: &[f64],
-    rhs: &[f64],
-    x0: Option<&[f64]>,
-) -> Result<Certificate, SolveError> {
-    let cold = x0
-        .map(|_| bidiag.certify(&vec![0.0; x.len()], rhs))
-        .transpose()?;
-    let mut cert = bidiag.certify(x, rhs)?;
-    if let Some(cold) = &cold {
-        cert.rebase(cold);
+/// Restarts from a refuted tolerance stop before the solve is refused outright.
+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);
+
+impl NormalEqReference {
+    /// A cold stream's first `(α₁, β₁)` is `‖Aᵀb‖` already.
+    fn cold(step1: BidiagStep) -> Self {
+        Self((step1.alpha * step1.beta).abs().max(f64::MIN_POSITIVE))
+    }
+
+    /// 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() {
+            Self(metric)
+        } else {
+            Self::cold(step1)
+        }
+    }
+
+    fn relative(self, estimate: f64) -> f64 {
+        estimate / self.0
     }
-    Ok(cert)
 }
 
-/// Runs the LSMR recurrences over a preconditioner-specific bidiagonalization stream.
-///
-/// The stream solves for the correction to `x0`; results and the audit are in terms of `b`.
+/// A warm start and its `‖Aᵀb‖`, since the stream's `(α₁, β₁)` measures `b − A x₀` instead.
+struct WarmStart<'a> {
+    x0: &'a [f64],
+    reference: NormalEqReference,
+}
+
+/// Runs LSMR on the correction to `x0`, restarting from any tolerance stop `b − A x` refutes.
 fn lsmr_from_bidiag<B: Bidiagonalization>(
     mut bidiag: B,
-    step1: BidiagStep,
+    mut step1: BidiagStep,
     b: &[f64],
-    x0: Option<&[f64]>,
+    warm_start: Option<WarmStart<'_>>,
     criteria: ConvergenceCriteria,
     maxiter: usize,
-    mut escalation: Option<Box<dyn EscalationHandler>>,
+    escalation: Option<&dyn EscalationPolicy>,
 ) -> Result<LsmrResult, SolveError> {
     let n = bidiag.v().len();
-    let total = |mut x: Vec<f64>| {
-        if let Some(x0) = x0 {
-            for (xi, &x0i) in x.iter_mut().zip(x0) {
-                *xi += x0i;
+    let reference = warm_start
+        .as_ref()
+        .map_or_else(|| NormalEqReference::cold(step1), |w| w.reference);
+    let mut base: Option<Cow<'_, [f64]>> = warm_start.map(|w| Cow::Borrowed(w.x0));
+    let mut iterations = 0;
+    let mut restarts = 0;
+    loop {
+        // A metric reporting no gradient at all may be hiding one outside itself.
+        if step1.alpha == 0.0 {
+            let x = base.map_or_else(|| vec![0.0; n], Cow::into_owned);
+            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 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 },
+                    )
+                }
+            };
+            return Ok(LsmrResult {
+                x,
+                converged,
+                iterations,
+                residual_norm: step1.beta,
+                normal_eq_residual,
+                stop_reason: if converged {
+                    LsmrStopReason::InitialNormalEquationResidualZero
+                } else {
+                    LsmrStopReason::FalseConvergence
+                },
+            });
+        }
+
+        // A pass reports its drop from its own ζ̄₀, so a restarted one needs a handler that agrees.
+        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 prev_rot = RotationStep::initial();
+        let stop_reason = 'pass: {
+            while iterations < maxiter {
+                iterations += 1;
+                let step = bidiag.step()?;
+                convergence.observe(step);
+                let curr_rot = recurrence.step(step);
+                solution.update(bidiag.v(), curr_rot, prev_rot);
+
+                // The tolerance test catches breakdown when the residual recurrences collapse.
+                match convergence.check(&recurrence) {
+                    Stop::Continue => {}
+                    Stop::ResidualTolerance => break 'pass LsmrStopReason::ResidualTolerance,
+                    Stop::NormalEquationTolerance => {
+                        break 'pass LsmrStopReason::NormalEquationTolerance
+                    }
+                }
+                if let Some(rule) = escalation.as_deref_mut() {
+                    let progress = Progress {
+                        iteration: iterations,
+                        normal_eq_residual: recurrence.relative_normal_eq_residual(),
+                    };
+                    if rule.should_escalate(progress) {
+                        break 'pass LsmrStopReason::Escalated;
+                    }
+                }
+                prev_rot = curr_rot;
             }
+            LsmrStopReason::MaxIterations
+        };
+
+        let mut x = solution.into_x();
+        if let Some(base) = &base {
+            axpby(&mut x, base, 1.0, 1.0);
         }
-        x
-    };
-    // A metric reporting no gradient at all may be hiding one outside itself.
-    if step1.alpha == 0.0 {
-        let x = total(vec![0.0; n]);
-        let cert = bidiag.certify(&x, b)?;
-        let converged = criteria.corroborated(&cert, || bidiag.operator_norm_below(b))?;
-        let (residual_norm, normal_eq_residual) = cert.residuals();
-        return Ok(LsmrResult {
+        let mut result = LsmrResult {
             x,
-            converged,
-            iterations: 0,
-            residual_norm,
-            normal_eq_residual,
-            stop_reason: if converged {
-                LsmrStopReason::InitialNormalEquationResidualZero
-            } else {
-                LsmrStopReason::FalseConvergence
-            },
-        });
-    }
-
-    let mut convergence = criteria.start(step1.alpha);
-    let mut recurrence = LsmrRecurrenceState::init(step1);
-    let mut solution = SolutionState::init(bidiag.v());
-    let mut prev_rot = RotationStep::initial();
-
-    let (iterations, stop_reason) = 'run: {
-        for itn in 1..=maxiter {
-            let step = bidiag.step()?;
-            convergence.observe(step);
-            let curr_rot = recurrence.step(step);
-            solution.update(bidiag.v(), curr_rot, prev_rot);
-
-            // The tolerance test catches breakdown when the residual recurrences collapse.
-            if let Some(stop_reason) = match convergence.check(&recurrence) {
-                Stop::Continue => None,
-                Stop::ResidualTolerance => Some(LsmrStopReason::ResidualTolerance),
-                Stop::NormalEquationTolerance => Some(LsmrStopReason::NormalEquationTolerance),
-            } {
-                let x = total(solution.into_x());
-                let cert = audit_against_rhs(&mut bidiag, &x, b, x0)?;
-                let converged = convergence.certified(&cert);
-                let (residual_norm, normal_eq_residual) = cert.residuals();
-                return Ok(LsmrResult {
-                    x,
-                    converged,
-                    iterations: itn,
-                    residual_norm,
-                    normal_eq_residual,
-                    stop_reason: if converged {
-                        stop_reason
-                    } else {
-                        LsmrStopReason::FalseConvergence
-                    },
-                });
+            converged: false,
+            iterations,
+            residual_norm: recurrence.residual_estimate(),
+            normal_eq_residual: reference.relative(recurrence.normal_eq_residual_estimate()),
+            stop_reason,
+        };
+        // Only a tolerance stop claims convergence, so only it is worth a true-residual evaluation.
+        if matches!(
+            stop_reason,
+            LsmrStopReason::ResidualTolerance | LsmrStopReason::NormalEquationTolerance
+        ) {
+            let residual_norm = bidiag.residual_norm(&result.x, b)?;
+            result.converged = criteria.corroborates_residual(residual_norm, result.residual_norm);
+            if !result.converged
+                && residual_norm.is_finite()
+                && restarts < MAX_RESTARTS
+                && iterations < maxiter
+            {
+                restarts += 1;
+                step1 = bidiag.restart(residual_norm)?;
+                base = Some(Cow::Owned(result.x));
+                continue;
             }
-            if let Some(rule) = escalation.as_deref_mut() {
-                let progress = Progress {
-                    iteration: itn,
-                    normal_eq_residual: recurrence.relative_normal_eq_residual(),
+            if !result.converged {
+                result.stop_reason = LsmrStopReason::FalseConvergence;
+                // 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)
+                } else {
+                    residual_norm
                 };
-                if rule.should_escalate(progress) {
-                    break 'run (itn, LsmrStopReason::Escalated);
-                }
             }
-            prev_rot = curr_rot;
+            result.residual_norm = residual_norm;
         }
-        (maxiter, LsmrStopReason::MaxIterations)
-    };
-    let x = total(solution.into_x());
-    // A warm run's recurrence measures the restart, so only an audit reports against `b`.
-    let (residual_norm, normal_eq_residual) = match x0 {
-        Some(_) => audit_against_rhs(&mut bidiag, &x, b, x0)?.residuals(),
-        None => (
-            recurrence.residual_estimate(),
-            recurrence.relative_normal_eq_residual(),
-        ),
-    };
-    Ok(LsmrResult {
-        x,
-        converged: false,
-        iterations,
-        residual_norm,
-        normal_eq_residual,
-        stop_reason,
-    })
+        return Ok(result);
+    }
 }
 
 fn validate_lsmr_inputs<A: Operator + ?Sized>(