Skip to main content

rusolver/
iterative.rs

1// SPDX-License-Identifier: Apache-2.0
2use crate::{
3    Matrix, MatrixView, SolverError, Tolerance, l2_norm
4};
5use crate::numerics::{
6    checked, dot, finite, sum_iter, zeros
7};
8/// Synchronous HOST linear map. Implementations must fully overwrite output,
9/// preserve inputs and implement a fixed linear map throughout one solve.
10/// Device-resident asynchronous algorithms will use a separate API, not hidden copies.
11pub trait LinearOperator {
12    fn dimension(&self)->usize;
13    fn apply(&self, x: &[f64], output: &mut[f64])->Result<(), SolverError>;
14}
15pub trait Preconditioner {
16    fn dimension(&self)->usize;
17    fn apply(&self, r: &[f64], output: &mut[f64])->Result<(), SolverError>;
18}
19impl LinearOperator for Matrix {
20    fn dimension(&self)->usize {
21        self.rows
22    }
23    fn apply(&self, x: &[f64], out: &mut[f64])->Result<(), SolverError> {
24        self.view().apply(x, out)
25    }
26}
27impl LinearOperator for MatrixView<'_> {
28    fn dimension(&self)->usize {
29        self.rows()
30    }
31    fn apply(&self, x: &[f64], out: &mut[f64])->Result<(), SolverError> {
32        let n=self.square()?;
33        if x.len()!=n || out.len()!=n {
34            return Err(SolverError::Shape("operator vector length"));
35        }
36        finite(x)?;
37        for i in 0..n {
38            out[i]=sum_iter((0..n).map(|j|self.at(i, j)*x[j]))?;
39        }
40        Ok(())
41    }
42}
43#[derive(Clone, Copy, Debug)]
44pub struct IdentityPreconditioner {
45    n: usize
46}
47impl IdentityPreconditioner {
48    pub fn new(n: usize)->Result<Self, SolverError> {
49        if n==0 {
50            Err(SolverError::Shape("empty preconditioner"))
51        } else {
52            Ok(Self{
53                n
54            })
55        }
56    }
57}
58impl Preconditioner for IdentityPreconditioner {
59    fn dimension(&self)->usize {
60        self.n
61    }
62    fn apply(&self, r: &[f64], out: &mut[f64])->Result<(), SolverError> {
63        if r.len()!=self.n || out.len()!=self.n {
64            return Err(SolverError::Shape("identity preconditioner"));
65        }
66        out.copy_from_slice(r);
67        Ok(())
68    }
69}
70/// Positive diagonal preconditioner. Reciprocal overflow is rejected at construction.
71#[derive(Clone, Debug)]
72pub struct JacobiPreconditioner {
73    inverse: Vec<f64>
74}
75impl JacobiPreconditioner {
76    pub fn new(diagonal: &[f64])->Result<Self, SolverError> {
77        if diagonal.is_empty() {
78            return Err(SolverError::Shape("empty Jacobi diagonal"));
79        }
80        finite(diagonal)?;
81        let mut inverse=zeros(diagonal.len())?;
82        for (i, &d) in diagonal.iter().enumerate() {
83            if d<=0.0 {
84                return Err(SolverError::NotPositiveDefinite{
85                    index: i, pivot: d
86                });
87            }
88            inverse[i]=checked(1.0/d, "Jacobi inverse diagonal")?;
89        }
90        Ok(Self{
91            inverse
92        })
93    }
94    pub fn from_matrix(a: MatrixView<'_>)->Result<Self, SolverError> {
95        let n=a.square()?;
96        let mut d=zeros(n)?;
97        for i in 0..n {
98            d[i]=a.at(i, i);
99        }
100        Self::new(&d)
101    }
102}
103impl Preconditioner for JacobiPreconditioner {
104    fn dimension(&self)->usize {
105        self.inverse.len()
106    }
107    fn apply(&self, r: &[f64], out: &mut[f64])->Result<(), SolverError> {
108        if r.len()!=self.inverse.len() || out.len()!=r.len() {
109            return Err(SolverError::Shape("Jacobi vector length"));
110        }
111        for i in 0..r.len() {
112            out[i]=checked(r[i]*self.inverse[i], "Jacobi application")?;
113        }
114        Ok(())
115    }
116}
117#[derive(Clone, Copy, Debug)]
118pub struct CgOptions {
119    pub tolerance: Tolerance, pub max_iterations: usize,
120    /// A true residual b-Ax is recomputed at this interval and before success.
121    /// Replacement restarts the search direction, trading some speed for reliability.
122    pub residual_recompute_interval: usize,
123}
124impl Default for CgOptions {
125    fn default()->Self {
126        Self{
127            tolerance: Tolerance{
128                absolute: 0.0, relative: 1e-10
129            }, max_iterations: 1000, residual_recompute_interval: 32
130        }
131    }
132}
133#[derive(Clone, Copy, Debug, PartialEq, Eq)]
134pub enum IterativeStatus {
135    Converged, MaxIterations
136}
137#[derive(Clone, Debug)]
138pub struct CgReport {
139    pub solution: Vec<f64>, pub iterations: usize, pub residual_norm: f64,
140    pub rhs_norm: f64, pub status: IterativeStatus, pub matvecs: usize
141}
142impl CgReport {
143    pub fn converged(&self)->bool {
144        self.status==IterativeStatus::Converged
145    }
146}
147fn checked_apply(a: &impl LinearOperator, x: &[f64], out: &mut[f64])->Result<(), SolverError> {
148    out.fill(f64::NAN);
149    a.apply(x, out)?;
150    finite(out)
151}
152fn checked_precondition(m: &impl Preconditioner, r: &[f64], out: &mut[f64])->Result<(), SolverError> {
153    out.fill(f64::NAN);
154    m.apply(r, out)?;
155    finite(out)
156}
157fn true_residual(a: &impl LinearOperator, b: &[f64], x: &[f64], r: &mut[f64], scratch: &mut[f64])->Result<f64, SolverError> {
158    checked_apply(a, x, scratch)?;
159    for i in 0..r.len() {
160        r[i]=checked(b[i]-scratch[i], "CG true residual")?;
161    }
162    l2_norm(r)
163}
164/// Preconditioned conjugate gradient for a fixed real symmetric POSITIVE-DEFINITE
165/// operator and positive-definite preconditioner. Caller guarantees those global
166/// properties; positive-curvature checks can detect some violations, not prove SPD.
167/// Zero RHS and an already-correct initial guess converge in zero iterations.
168/// Failure to reach tolerance returns MaxIterations, never a false success.
169pub fn conjugate_gradient(a: &impl LinearOperator, b: &[f64], initial: Option<&[f64]>,
170preconditioner: &impl Preconditioner, options: CgOptions)->Result<CgReport, SolverError> {
171    let n=a.dimension();
172    if n==0 || b.len()!=n || preconditioner.dimension()!=n {
173        return Err(SolverError::Shape("CG dimension"));
174    }
175    options.tolerance.validate()?;
176    if options.residual_recompute_interval==0 || (options.tolerance.absolute==0.0 && options.tolerance.relative==0.0) {
177        return Err(SolverError::InvalidOption("CG tolerance and residual interval must be positive"));
178    }
179    finite(b)?;
180    let mut x=zeros(n)?;
181    if let Some(initial)=initial {
182        if initial.len()!=n {
183            return Err(SolverError::Shape("CG initial guess"));
184        }
185        finite(initial)?;
186        x.copy_from_slice(initial);
187    }
188    let (mut r, mut z, mut p, mut ap)=(zeros(n)?, zeros(n)?, zeros(n)?, zeros(n)?);
189    let rhs_norm=l2_norm(b)?;
190    let target=options.tolerance.threshold(rhs_norm)?;
191    let mut matvecs=1;
192    let mut residual_norm=true_residual(a, b, &x, &mut r, &mut ap)?;
193    let report=|solution, iterations, residual_norm, status, matvecs|CgReport{
194        solution, iterations, residual_norm, rhs_norm, status, matvecs
195    };
196    if residual_norm<=target {
197        return Ok(report(x, 0, residual_norm, IterativeStatus::Converged, matvecs));
198    }
199    if options.max_iterations==0 {
200        return Ok(report(x, 0, residual_norm, IterativeStatus::MaxIterations, matvecs));
201    }
202    checked_precondition(preconditioner, &r, &mut z)?;
203    let mut rho=dot(&r, &z)?;
204    if rho<=0.0 {
205        return Err(SolverError::Breakdown("nonpositive r^T M^-1 r (or numerical underflow)"));
206    }
207    p.copy_from_slice(&z);
208    for iteration in 1..=options.max_iterations {
209        checked_apply(a, &p, &mut ap)?;
210        matvecs+=1;
211        let curvature=dot(&p, &ap)?;
212        if curvature<=0.0 {
213            return Err(SolverError::Breakdown("nonpositive p^T A p; SPD required"));
214        }
215        let alpha=checked(rho/curvature, "CG step size")?;
216        for i in 0..n {
217            x[i]=checked(alpha.mul_add(p[i], x[i]), "CG iterate")?;
218            r[i]=checked((-alpha).mul_add(ap[i], r[i]), "CG recursive residual")?;
219        }
220        residual_norm=l2_norm(&r)?;
221        let replace=iteration%options.residual_recompute_interval==0 || residual_norm<=target || iteration==options.max_iterations;
222        if replace {
223            residual_norm=true_residual(a, b, &x, &mut r, &mut ap)?;
224            matvecs+=1;
225        }
226        if residual_norm<=target {
227            return Ok(report(x, iteration, residual_norm, IterativeStatus::Converged, matvecs));
228        }
229        if iteration==options.max_iterations {
230            break;
231        }
232        checked_precondition(preconditioner, &r, &mut z)?;
233        let next_rho=dot(&r, &z)?;
234        if next_rho<=0.0 {
235            return Err(SolverError::Breakdown("preconditioned residual underflow/nonpositive"));
236        }
237        if replace {
238            p.copy_from_slice(&z);
239        } else {
240            let beta=checked(next_rho/rho, "CG direction coefficient")?;
241            for i in 0..n {
242                p[i]=checked(beta.mul_add(p[i], z[i]), "CG direction")?;
243            }
244        }
245        rho=next_rho;
246    }
247    Ok(report(x, options.max_iterations, residual_norm, IterativeStatus::MaxIterations, matvecs))
248}