use crate::error::Failed;
use crate::linalg::basic::arrays::{Array, Array1, Array2, ArrayView1, MutArrayView1};
use crate::numbers::floatnum::FloatNumber;
pub trait BiconjugateGradientSolver<'a, T: FloatNumber, X: Array2<T>> {
fn solve_mut(
&self,
a: &'a X,
b: &Vec<T>,
x: &mut Vec<T>,
tol: T,
max_iter: usize,
) -> Result<T, Failed> {
if tol <= T::zero() {
return Err(Failed::fit("tolerance shoud be > 0"));
}
if max_iter == 0 {
return Err(Failed::fit("maximum number of iterations should be > 0"));
}
let n = b.shape();
let mut r = Vec::zeros(n);
let mut rr = Vec::zeros(n);
let mut z = Vec::zeros(n);
let mut zz = Vec::zeros(n);
self.mat_vec_mul(a, x, &mut r);
for j in 0..n {
r[j] = b[j] - r[j];
rr[j] = r[j];
}
let bnrm = b.norm(2f64);
self.solve_preconditioner(a, &r[..], &mut z[..]);
let mut p = Vec::zeros(n);
let mut pp = Vec::zeros(n);
let mut bkden = T::zero();
let mut err = T::zero();
for iter in 1..max_iter {
let mut bknum = T::zero();
self.solve_preconditioner(a, &rr, &mut zz);
for j in 0..n {
bknum += z[j] * rr[j];
}
if iter == 1 {
p[..n].copy_from_slice(&z[..n]);
pp[..n].copy_from_slice(&zz[..n]);
} else {
let bk = bknum / bkden;
for j in 0..n {
p[j] = bk * pp[j] + z[j];
pp[j] = bk * pp[j] + zz[j];
}
}
bkden = bknum;
self.mat_vec_mul(a, &p, &mut z);
let mut akden = T::zero();
for j in 0..n {
akden += z[j] * pp[j];
}
let ak = bknum / akden;
self.mat_t_vec_mul(a, &pp, &mut zz);
for j in 0..n {
x[j] += ak * p[j];
r[j] -= ak * z[j];
rr[j] -= ak * zz[j];
}
self.solve_preconditioner(a, &r, &mut z);
err = T::from_f64(r.norm(2f64) / bnrm).unwrap();
if err <= tol {
break;
}
}
Ok(err)
}
fn solve_preconditioner(&self, a: &'a X, b: &[T], x: &mut [T]) {
let diag = Self::diag(a);
let n = diag.len();
for (i, diag_i) in diag.iter().enumerate().take(n) {
if *diag_i != T::zero() {
x[i] = b[i] / *diag_i;
} else {
x[i] = b[i];
}
}
}
fn mat_vec_mul(&self, a: &X, x: &Vec<T>, y: &mut Vec<T>) {
y.copy_from(&x.xa(false, a));
}
fn mat_t_vec_mul(&self, a: &X, x: &Vec<T>, y: &mut Vec<T>) {
y.copy_from(&x.xa(true, a));
}
fn diag(a: &X) -> Vec<T> {
let (nrows, ncols) = a.shape();
let n = nrows.min(ncols);
let mut d = Vec::with_capacity(n);
for i in 0..n {
d.push(*a.get((i, i)));
}
d
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::linalg::basic::arrays::Array2;
use crate::linalg::basic::matrix::DenseMatrix;
pub struct BGSolver {}
impl<T: FloatNumber, X: Array2<T>> BiconjugateGradientSolver<'_, T, X> for BGSolver {}
#[test]
fn bg_solver() {
let a = DenseMatrix::from_2d_array(&[&[25., 15., -5.], &[15., 18., 0.], &[-5., 0., 11.]]);
let b = vec![40., 51., 28.];
let expected = vec![1.0, 2.0, 3.0];
let mut x = Vec::zeros(3);
let solver = BGSolver {};
let err: f64 = solver.solve_mut(&a, &b, &mut x, 1e-6, 6).unwrap();
assert!(x
.iter()
.zip(expected.iter())
.all(|(&a, &b)| (a - b).abs() < 1e-4));
assert!((err - 0.0).abs() < 1e-4);
}
}