1use crate::{
3 Matrix, MatrixView, SolverError, Tolerance, l2_norm
4};
5use crate::numerics::{
6 checked, dot, finite, sum_iter, zeros
7};
8pub 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#[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 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}
164pub 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}