uncertain_numerics/
covariance_greedy_solver.rs

1use nalgebra::{DMatrix, DVector};
2
3use crate::{GaussianLinearBelief, LinearSolverError, SpdLinearSystem};
4
5const INFORMATION_TOLERANCE: f64 = 128.0 * f64::EPSILON;
6
7/// Candidate projection chosen to maximize posterior covariance-trace reduction.
8#[derive(Debug, Clone, PartialEq)]
9pub struct SelectedLinearDirection {
10    direction: Vec<f64>,
11    index: usize,
12    trace_reduction: f64,
13}
14
15impl SelectedLinearDirection {
16    /// Return the selected direction.
17    #[must_use]
18    pub fn direction(&self) -> &[f64] {
19        &self.direction
20    }
21
22    /// Return the original candidate index.
23    #[must_use]
24    pub const fn index(&self) -> usize {
25        self.index
26    }
27
28    /// Return the predicted posterior covariance-trace reduction.
29    #[must_use]
30    pub const fn trace_reduction(&self) -> f64 {
31        self.trace_reduction
32    }
33}
34
35/// Reason a covariance-greedy probabilistic solve terminated.
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub enum CovarianceGreedyTermination {
38    /// Posterior covariance trace reached the requested tolerance.
39    CovarianceTraceToleranceReached,
40    /// The configured projection budget was exhausted.
41    ProjectionBudgetReached,
42    /// No unevaluated candidate directions remain.
43    CandidatesExhausted,
44    /// Remaining candidates carry no posterior uncertainty.
45    NoInformativeDirection,
46}
47
48/// One covariance-greedy projection step.
49#[derive(Debug, Clone, PartialEq)]
50pub struct CovarianceGreedyStep {
51    direction: Vec<f64>,
52    predicted_trace_reduction: f64,
53    posterior_trace: f64,
54}
55
56impl CovarianceGreedyStep {
57    /// Return the selected direction.
58    #[must_use]
59    pub fn direction(&self) -> &[f64] {
60        &self.direction
61    }
62
63    /// Return the predicted covariance-trace reduction before conditioning.
64    #[must_use]
65    pub const fn predicted_trace_reduction(&self) -> f64 {
66        self.predicted_trace_reduction
67    }
68
69    /// Return the posterior covariance trace after conditioning.
70    #[must_use]
71    pub const fn posterior_trace(&self) -> f64 {
72        self.posterior_trace
73    }
74}
75
76/// Result of a covariance-greedy probabilistic linear solve.
77#[derive(Debug, Clone, PartialEq)]
78pub struct CovarianceGreedySolveResult {
79    belief: GaussianLinearBelief,
80    steps: Vec<CovarianceGreedyStep>,
81    remaining_candidates: Vec<Vec<f64>>,
82    termination: CovarianceGreedyTermination,
83}
84
85impl CovarianceGreedySolveResult {
86    /// Return the final Gaussian solution belief.
87    #[must_use]
88    pub const fn belief(&self) -> &GaussianLinearBelief {
89        &self.belief
90    }
91
92    /// Return the ordered projection history.
93    #[must_use]
94    pub fn steps(&self) -> &[CovarianceGreedyStep] {
95        &self.steps
96    }
97
98    /// Return candidate directions not selected during the run.
99    #[must_use]
100    pub fn remaining_candidates(&self) -> &[Vec<f64>] {
101        &self.remaining_candidates
102    }
103
104    /// Return the termination reason.
105    #[must_use]
106    pub const fn termination(&self) -> CovarianceGreedyTermination {
107        self.termination
108    }
109}
110
111/// Covariance-trace acquisition for exact linear-system projection observations.
112#[derive(Debug, Clone, Copy, PartialEq)]
113pub struct CovarianceTraceAcquisition;
114
115impl CovarianceTraceAcquisition {
116    /// Predict the covariance-trace reduction from conditioning on one direction.
117    ///
118    /// For `h = A^T s`, the exact rank-one covariance update implies
119    ///
120    /// ```text
121    /// trace(Sigma) - trace(Sigma+)
122    /// = h^T Sigma^2 h / (h^T Sigma h).
123    /// ```
124    ///
125    /// # Errors
126    ///
127    /// Returns [`LinearSolverError`] for incompatible/non-finite directions.
128    pub fn reduction(
129        system: &SpdLinearSystem,
130        belief: &GaussianLinearBelief,
131        direction: &[f64],
132    ) -> Result<f64, LinearSolverError> {
133        if system.dimension() != belief.dimension() || direction.len() != system.dimension() {
134            return Err(LinearSolverError::VectorDimensionMismatch);
135        }
136        if direction.iter().any(|value| !value.is_finite()) {
137            return Err(LinearSolverError::NonFiniteVectorEntry);
138        }
139
140        let dimension = system.dimension();
141        let matrix = DMatrix::from_row_slice(dimension, dimension, system.matrix());
142        let covariance = DMatrix::from_row_slice(dimension, dimension, belief.covariance());
143        let search = DVector::from_column_slice(direction);
144        let h = matrix.transpose() * search;
145        let covariance_h = &covariance * &h;
146        let denominator = h.dot(&covariance_h);
147        let scale = covariance
148            .iter()
149            .fold(1.0_f64, |acc, value| acc.max(value.abs()));
150        let tolerance = INFORMATION_TOLERANCE * scale;
151        if denominator <= tolerance {
152            return Ok(0.0);
153        }
154
155        Ok(covariance_h.dot(&covariance_h) / denominator)
156    }
157
158    /// Select the direction with the largest predicted covariance-trace reduction.
159    ///
160    /// Ties are resolved deterministically by retaining the first maximum.
161    ///
162    /// # Errors
163    ///
164    /// Returns [`LinearSolverError`] for an empty candidate set or invalid directions.
165    pub fn select_best(
166        system: &SpdLinearSystem,
167        belief: &GaussianLinearBelief,
168        candidates: &[Vec<f64>],
169    ) -> Result<SelectedLinearDirection, LinearSolverError> {
170        if candidates.is_empty() {
171            return Err(LinearSolverError::EmptyCandidateDirections);
172        }
173
174        let mut best = SelectedLinearDirection {
175            direction: candidates[0].clone(),
176            index: 0,
177            trace_reduction: Self::reduction(system, belief, &candidates[0])?,
178        };
179        for (index, candidate) in candidates.iter().enumerate().skip(1) {
180            let reduction = Self::reduction(system, belief, candidate)?;
181            if reduction > best.trace_reduction {
182                best = SelectedLinearDirection {
183                    direction: candidate.clone(),
184                    index,
185                    trace_reduction: reduction,
186                };
187            }
188        }
189        Ok(best)
190    }
191}
192
193/// Sequential probabilistic linear solver with data-independent covariance-greedy directions.
194#[derive(Debug, Clone, Copy, PartialEq)]
195pub struct CovarianceGreedyProjectionSolver {
196    covariance_trace_tolerance: f64,
197    max_projections: usize,
198}
199
200impl CovarianceGreedyProjectionSolver {
201    /// Construct the solver with explicit uncertainty tolerance and projection budget.
202    ///
203    /// # Errors
204    ///
205    /// Returns [`LinearSolverError`] when the covariance-trace tolerance is invalid.
206    pub fn new(
207        covariance_trace_tolerance: f64,
208        max_projections: usize,
209    ) -> Result<Self, LinearSolverError> {
210        if !covariance_trace_tolerance.is_finite() {
211            return Err(LinearSolverError::NonFiniteTolerance);
212        }
213        if covariance_trace_tolerance < 0.0 {
214            return Err(LinearSolverError::NegativeTolerance);
215        }
216        Ok(Self {
217            covariance_trace_tolerance,
218            max_projections,
219        })
220    }
221
222    /// Run covariance-greedy sequential conditioning from an initial Gaussian belief.
223    ///
224    /// Direction selection and stopping depend only on `A`, the candidate set,
225    /// the covariance state, and the configured budget/tolerance. They do not
226    /// depend on the observed right-hand side.
227    ///
228    /// # Errors
229    ///
230    /// Returns [`LinearSolverError`] for incompatible dimensions or invalid candidates.
231    pub fn solve(
232        &self,
233        system: &SpdLinearSystem,
234        initial_belief: &GaussianLinearBelief,
235        candidates: &[Vec<f64>],
236    ) -> Result<CovarianceGreedySolveResult, LinearSolverError> {
237        if system.dimension() != initial_belief.dimension() {
238            return Err(LinearSolverError::VectorDimensionMismatch);
239        }
240        if candidates.is_empty() {
241            return Err(LinearSolverError::EmptyCandidateDirections);
242        }
243
244        let mut belief = initial_belief.clone();
245        let mut remaining_candidates = candidates.to_vec();
246        let mut steps = Vec::new();
247
248        if covariance_trace(&belief) <= self.covariance_trace_tolerance {
249            return Ok(CovarianceGreedySolveResult {
250                belief,
251                steps,
252                remaining_candidates,
253                termination: CovarianceGreedyTermination::CovarianceTraceToleranceReached,
254            });
255        }
256
257        for _ in 0..self.max_projections {
258            if remaining_candidates.is_empty() {
259                return Ok(CovarianceGreedySolveResult {
260                    belief,
261                    steps,
262                    remaining_candidates,
263                    termination: CovarianceGreedyTermination::CandidatesExhausted,
264                });
265            }
266
267            let selected =
268                CovarianceTraceAcquisition::select_best(system, &belief, &remaining_candidates)?;
269            if selected.trace_reduction() <= INFORMATION_TOLERANCE {
270                return Ok(CovarianceGreedySolveResult {
271                    belief,
272                    steps,
273                    remaining_candidates,
274                    termination: CovarianceGreedyTermination::NoInformativeDirection,
275                });
276            }
277
278            let trace_before = covariance_trace(&belief);
279            let updated = belief.condition_on_projection(system, selected.direction())?;
280            let trace_after = covariance_trace(&updated);
281            let actual_reduction = trace_before - trace_after;
282            let tolerance = 1.0e-10 * selected.trace_reduction().abs().max(1.0);
283            debug_assert!((actual_reduction - selected.trace_reduction()).abs() <= tolerance);
284
285            remaining_candidates.remove(selected.index());
286            steps.push(CovarianceGreedyStep {
287                direction: selected.direction,
288                predicted_trace_reduction: selected.trace_reduction,
289                posterior_trace: trace_after,
290            });
291            belief = updated;
292
293            if trace_after <= self.covariance_trace_tolerance {
294                return Ok(CovarianceGreedySolveResult {
295                    belief,
296                    steps,
297                    remaining_candidates,
298                    termination: CovarianceGreedyTermination::CovarianceTraceToleranceReached,
299                });
300            }
301        }
302
303        let termination = if remaining_candidates.is_empty() {
304            CovarianceGreedyTermination::CandidatesExhausted
305        } else {
306            CovarianceGreedyTermination::ProjectionBudgetReached
307        };
308        Ok(CovarianceGreedySolveResult {
309            belief,
310            steps,
311            remaining_candidates,
312            termination,
313        })
314    }
315}
316
317fn covariance_trace(belief: &GaussianLinearBelief) -> f64 {
318    let dimension = belief.dimension();
319    (0..dimension)
320        .map(|index| belief.covariance()[index * dimension + index])
321        .sum()
322}
323
324#[cfg(test)]
325mod tests {
326    use super::{
327        CovarianceGreedyProjectionSolver, CovarianceGreedyTermination, CovarianceTraceAcquisition,
328    };
329    use crate::{GaussianLinearBelief, LinearSolverError, SpdLinearSystem};
330
331    fn identity_belief(dimension: usize) -> GaussianLinearBelief {
332        let mut covariance = vec![0.0; dimension * dimension];
333        for index in 0..dimension {
334            covariance[index * dimension + index] = 1.0;
335        }
336        GaussianLinearBelief::new(&vec![0.0; dimension], &covariance, dimension)
337            .expect("identity belief is valid")
338    }
339
340    #[test]
341    fn predicted_trace_reduction_matches_actual_conditioning_drop() {
342        let system =
343            SpdLinearSystem::new(&[4.0, 1.0, 1.0, 3.0], &[1.0, 2.0], 2).expect("system is valid");
344        let belief = identity_belief(2);
345        let direction = [1.0, 0.0];
346        let predicted = CovarianceTraceAcquisition::reduction(&system, &belief, &direction)
347            .expect("direction is valid");
348        let updated = belief
349            .condition_on_projection(&system, &direction)
350            .expect("direction is informative");
351        let actual = 2.0 - updated.covariance()[0] - updated.covariance()[3];
352        assert!((predicted - actual).abs() <= 1.0e-12);
353    }
354
355    #[test]
356    fn selector_returns_global_trace_reduction_maximum() {
357        let system =
358            SpdLinearSystem::new(&[4.0, 1.0, 1.0, 3.0], &[1.0, 2.0], 2).expect("system is valid");
359        let belief = identity_belief(2);
360        let candidates = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]];
361        let selected = CovarianceTraceAcquisition::select_best(&system, &belief, &candidates)
362            .expect("candidates are valid");
363        for candidate in &candidates {
364            let reduction = CovarianceTraceAcquisition::reduction(&system, &belief, candidate)
365                .expect("candidate is valid");
366            assert!(selected.trace_reduction() + 1.0e-12 >= reduction);
367        }
368    }
369
370    #[test]
371    fn sequential_selection_does_not_use_rhs_values() {
372        let matrix = [4.0, 1.0, 1.0, 3.0];
373        let first_system = SpdLinearSystem::new(&matrix, &[1.0, 2.0], 2).expect("valid system");
374        let second_system = SpdLinearSystem::new(&matrix, &[-7.0, 4.0], 2).expect("valid system");
375        let candidates = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]];
376        let solver = CovarianceGreedyProjectionSolver::new(0.0, 2).expect("solver is valid");
377        let first = solver
378            .solve(&first_system, &identity_belief(2), &candidates)
379            .expect("solve is valid");
380        let second = solver
381            .solve(&second_system, &identity_belief(2), &candidates)
382            .expect("solve is valid");
383        let first_directions: Vec<&[f64]> = first
384            .steps()
385            .iter()
386            .map(super::CovarianceGreedyStep::direction)
387            .collect();
388        let second_directions: Vec<&[f64]> = second
389            .steps()
390            .iter()
391            .map(super::CovarianceGreedyStep::direction)
392            .collect();
393        assert_eq!(first_directions, second_directions);
394    }
395
396    #[test]
397    fn reports_projection_budget_termination() {
398        let system =
399            SpdLinearSystem::new(&[4.0, 1.0, 1.0, 3.0], &[1.0, 2.0], 2).expect("system is valid");
400        let candidates = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]];
401        let solver = CovarianceGreedyProjectionSolver::new(0.0, 1).expect("solver is valid");
402        let result = solver
403            .solve(&system, &identity_belief(2), &candidates)
404            .expect("solve is valid");
405        assert_eq!(
406            result.termination(),
407            CovarianceGreedyTermination::ProjectionBudgetReached
408        );
409        assert_eq!(result.steps().len(), 1);
410    }
411
412    #[test]
413    fn rejects_empty_candidates() {
414        let system = SpdLinearSystem::new(&[1.0], &[1.0], 1).expect("system is valid");
415        let solver = CovarianceGreedyProjectionSolver::new(0.0, 1).expect("solver is valid");
416        assert_eq!(
417            solver.solve(&system, &identity_belief(1), &[]),
418            Err(LinearSolverError::EmptyCandidateDirections)
419        );
420    }
421}
uncertain_numerics/covariance_greedy_solver.rs

uncertain_numerics/
covariance_greedy_solver.rs
