uncertain_numerics/
probabilistic_linear_solver.rs

1use nalgebra::{DMatrix, DVector};
2
3use crate::{GaussianLinearBelief, LinearSolverError, SpdLinearSystem};
4
5/// Reason a residual-projection probabilistic linear solve terminated.
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum LinearSolveTermination {
8    /// The residual norm reached the requested tolerance.
9    ResidualToleranceReached,
10    /// The total posterior variance reached the requested tolerance.
11    CovarianceTraceToleranceReached,
12    /// The configured iteration budget was exhausted.
13    IterationBudgetReached,
14    /// The current residual direction contains no remaining posterior uncertainty.
15    NoInformativeDirection,
16}
17
18/// One residual-projection conditioning step.
19#[derive(Debug, Clone, PartialEq)]
20pub struct LinearSolveStep {
21    residual_norm_before: f64,
22    residual_norm_after: f64,
23    covariance_trace_after: f64,
24    search_direction: Vec<f64>,
25}
26
27impl LinearSolveStep {
28    /// Residual norm `||b - A m||` of the belief mean before this projection.
29    #[must_use]
30    pub const fn residual_norm_before(&self) -> f64 {
31        self.residual_norm_before
32    }
33
34    /// Residual norm of the updated belief mean after this projection.
35    #[must_use]
36    pub const fn residual_norm_after(&self) -> f64 {
37        self.residual_norm_after
38    }
39
40    /// Trace of the posterior covariance after this projection.
41    #[must_use]
42    pub const fn covariance_trace_after(&self) -> f64 {
43        self.covariance_trace_after
44    }
45
46    /// Unit search direction `s` defining the projection `s^T A x = s^T b`.
47    #[must_use]
48    pub fn search_direction(&self) -> &[f64] {
49        &self.search_direction
50    }
51}
52
53/// Result of an iterative probabilistic linear solve.
54#[derive(Debug, Clone, PartialEq)]
55pub struct ProbabilisticLinearSolveResult {
56    belief: GaussianLinearBelief,
57    steps: Vec<LinearSolveStep>,
58    termination: LinearSolveTermination,
59}
60
61impl ProbabilisticLinearSolveResult {
62    /// Final Gaussian belief over the solution vector.
63    #[must_use]
64    pub const fn belief(&self) -> &GaussianLinearBelief {
65        &self.belief
66    }
67
68    /// Ordered history of conditioning steps.
69    #[must_use]
70    pub fn steps(&self) -> &[LinearSolveStep] {
71        &self.steps
72    }
73
74    /// Reason the solve stopped.
75    #[must_use]
76    pub const fn termination(&self) -> LinearSolveTermination {
77        self.termination
78    }
79}
80
81/// Residual-driven probabilistic solver for dense SPD systems.
82///
83/// At iteration `k`, the normalized residual of the current posterior mean,
84/// `s_k = r_k / ||r_k||`, defines the next exact projection observation
85/// `s_k^T A x = s_k^T b`.
86#[derive(Debug, Clone, Copy, PartialEq)]
87pub struct ResidualProjectionSolver {
88    residual_tolerance: f64,
89    covariance_trace_tolerance: f64,
90    max_iterations: usize,
91}
92
93impl ResidualProjectionSolver {
94    /// Construct the solver with explicit stopping tolerances.
95    ///
96    /// # Errors
97    ///
98    /// Returns [`LinearSolverError`] when either tolerance is non-finite or negative.
99    pub fn new(
100        residual_tolerance: f64,
101        covariance_trace_tolerance: f64,
102        max_iterations: usize,
103    ) -> Result<Self, LinearSolverError> {
104        if !residual_tolerance.is_finite() || !covariance_trace_tolerance.is_finite() {
105            return Err(LinearSolverError::NonFiniteTolerance);
106        }
107        if residual_tolerance < 0.0 || covariance_trace_tolerance < 0.0 {
108            return Err(LinearSolverError::NegativeTolerance);
109        }
110        Ok(Self {
111            residual_tolerance,
112            covariance_trace_tolerance,
113            max_iterations,
114        })
115    }
116
117    /// Solve from an initial Gaussian belief.
118    ///
119    /// # Errors
120    ///
121    /// Returns [`LinearSolverError`] when the belief dimension does not match the
122    /// system or when a conditioning step fails for a reason other than a
123    /// degenerate residual projection.
124    pub fn solve(
125        &self,
126        system: &SpdLinearSystem,
127        initial_belief: &GaussianLinearBelief,
128    ) -> Result<ProbabilisticLinearSolveResult, LinearSolverError> {
129        if system.dimension() != initial_belief.dimension() {
130            return Err(LinearSolverError::VectorDimensionMismatch);
131        }
132
133        let dimension = system.dimension();
134        let matrix = DMatrix::from_row_slice(dimension, dimension, system.matrix());
135        let rhs = DVector::from_column_slice(system.rhs());
136        let mut belief = initial_belief.clone();
137        let mut steps = Vec::new();
138
139        let mut residual = compute_residual(&matrix, &rhs, belief.mean());
140        let mut residual_norm = residual.norm();
141        if residual_norm <= self.residual_tolerance {
142            return Ok(ProbabilisticLinearSolveResult {
143                belief,
144                steps,
145                termination: LinearSolveTermination::ResidualToleranceReached,
146            });
147        }
148        if covariance_trace(&belief) <= self.covariance_trace_tolerance {
149            return Ok(ProbabilisticLinearSolveResult {
150                belief,
151                steps,
152                termination: LinearSolveTermination::CovarianceTraceToleranceReached,
153            });
154        }
155
156        for _ in 0..self.max_iterations {
157            let search = &residual / residual_norm;
158            let updated = match belief.condition_on_projection(system, search.as_slice()) {
159                Ok(updated) => updated,
160                Err(LinearSolverError::DegenerateObservation) => {
161                    return Ok(ProbabilisticLinearSolveResult {
162                        belief,
163                        steps,
164                        termination: LinearSolveTermination::NoInformativeDirection,
165                    });
166                }
167                Err(error) => return Err(error),
168            };
169
170            let next_residual = compute_residual(&matrix, &rhs, updated.mean());
171            let next_residual_norm = next_residual.norm();
172            let trace = covariance_trace(&updated);
173            steps.push(LinearSolveStep {
174                residual_norm_before: residual_norm,
175                residual_norm_after: next_residual_norm,
176                covariance_trace_after: trace,
177                search_direction: search.as_slice().to_vec(),
178            });
179            belief = updated;
180            residual = next_residual;
181            residual_norm = next_residual_norm;
182
183            if residual_norm <= self.residual_tolerance {
184                return Ok(ProbabilisticLinearSolveResult {
185                    belief,
186                    steps,
187                    termination: LinearSolveTermination::ResidualToleranceReached,
188                });
189            }
190            if trace <= self.covariance_trace_tolerance {
191                return Ok(ProbabilisticLinearSolveResult {
192                    belief,
193                    steps,
194                    termination: LinearSolveTermination::CovarianceTraceToleranceReached,
195                });
196            }
197        }
198
199        Ok(ProbabilisticLinearSolveResult {
200            belief,
201            steps,
202            termination: LinearSolveTermination::IterationBudgetReached,
203        })
204    }
205}
206
207fn compute_residual(matrix: &DMatrix<f64>, rhs: &DVector<f64>, mean: &[f64]) -> DVector<f64> {
208    rhs - matrix * DVector::from_column_slice(mean)
209}
210
211fn covariance_trace(belief: &GaussianLinearBelief) -> f64 {
212    let dimension = belief.dimension();
213    (0..dimension)
214        .map(|index| belief.covariance()[index * dimension + index])
215        .sum()
216}
217
218#[cfg(test)]
219mod tests {
220    use super::{LinearSolveTermination, ResidualProjectionSolver};
221    use crate::{GaussianLinearBelief, LinearSolverError, SpdLinearSystem};
222
223    fn identity_belief(dimension: usize) -> GaussianLinearBelief {
224        let mut covariance = vec![0.0; dimension * dimension];
225        for index in 0..dimension {
226            covariance[index * dimension + index] = 1.0;
227        }
228        GaussianLinearBelief::new(&vec![0.0; dimension], &covariance, dimension)
229            .expect("identity belief is valid")
230    }
231
232    #[test]
233    fn solves_two_dimensional_spd_system_in_at_most_dimension_steps() {
234        let system =
235            SpdLinearSystem::new(&[4.0, 1.0, 1.0, 3.0], &[1.0, 2.0], 2).expect("system is valid");
236        let solver = ResidualProjectionSolver::new(1.0e-12, 0.0, 2).expect("solver is valid");
237        let result = solver
238            .solve(&system, &identity_belief(2))
239            .expect("solve should succeed");
240
241        assert_eq!(
242            result.termination(),
243            LinearSolveTermination::ResidualToleranceReached
244        );
245        assert!(result.steps().len() <= 2);
246        assert!((result.belief().mean()[0] - 1.0 / 11.0).abs() < 1.0e-10);
247        assert!((result.belief().mean()[1] - 7.0 / 11.0).abs() < 1.0e-10);
248    }
249
250    #[test]
251    fn covariance_trace_is_non_increasing() {
252        let system = SpdLinearSystem::new(
253            &[1.0, 0.0, 0.0, 0.0, 3.0, 0.0, 0.0, 0.0, 10.0],
254            &[1.0, 2.0, -1.0],
255            3,
256        )
257        .expect("system is valid");
258        let solver = ResidualProjectionSolver::new(0.0, 0.0, 3).expect("solver is valid");
259        let result = solver
260            .solve(&system, &identity_belief(3))
261            .expect("solve should succeed");
262
263        let mut previous = 3.0;
264        for step in result.steps() {
265            assert!(step.covariance_trace_after() <= previous + 1.0e-12);
266            previous = step.covariance_trace_after();
267        }
268    }
269
270    #[test]
271    fn residual_norm_need_not_be_monotone() {
272        let system = SpdLinearSystem::new(
273            &[1.0, 0.0, 0.0, 0.0, 3.0, 0.0, 0.0, 0.0, 10.0],
274            &[1.0, 2.0, -1.0],
275            3,
276        )
277        .expect("system is valid");
278        let solver = ResidualProjectionSolver::new(0.0, 0.0, 3).expect("solver is valid");
279        let result = solver
280            .solve(&system, &identity_belief(3))
281            .expect("solve should succeed");
282
283        assert!(
284            result
285                .steps()
286                .iter()
287                .any(|step| step.residual_norm_after() > step.residual_norm_before())
288        );
289    }
290
291    #[test]
292    fn supports_covariance_trace_stopping() {
293        let system =
294            SpdLinearSystem::new(&[2.0, 0.0, 0.0, 1.0], &[1.0, -1.0], 2).expect("system is valid");
295        let solver = ResidualProjectionSolver::new(0.0, 1.1, 5).expect("solver is valid");
296        let result = solver
297            .solve(&system, &identity_belief(2))
298            .expect("solve should succeed");
299
300        assert_eq!(
301            result.termination(),
302            LinearSolveTermination::CovarianceTraceToleranceReached
303        );
304        assert_eq!(result.steps().len(), 1);
305    }
306
307    #[test]
308    fn validates_tolerances() {
309        assert_eq!(
310            ResidualProjectionSolver::new(f64::NAN, 0.0, 1),
311            Err(LinearSolverError::NonFiniteTolerance)
312        );
313        assert_eq!(
314            ResidualProjectionSolver::new(0.0, -1.0, 1),
315            Err(LinearSolverError::NegativeTolerance)
316        );
317    }
318}
uncertain_numerics/probabilistic_linear_solver.rs

uncertain_numerics/
probabilistic_linear_solver.rs
