feat: adds LASSO

src/linear/bg_solver.rs

@@ -0,0 +1,146 @@
+//! This is a generic solver for Ax = b type of equation
+//!
+//! for more information take a look at [this Wikipedia article](https://en.wikipedia.org/wiki/Biconjugate_gradient_method)
+//! and [this paper](https://www.cs.cmu.edu/~quake-papers/painless-conjugate-gradient.pdf)
+use crate::error::Failed;
+use crate::linalg::Matrix;
+use crate::math::num::RealNumber;
+
+pub trait BiconjugateGradientSolver<T: RealNumber, M: Matrix<T>> {
+    fn solve_mut(&self, a: &M, b: &M, x: &mut M, 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 = M::zeros(n, 1);
+        let mut rr = M::zeros(n, 1);
+        let mut z = M::zeros(n, 1);
+        let mut zz = M::zeros(n, 1);
+
+        self.mat_vec_mul(a, x, &mut r);
+
+        for j in 0..n {
+            r.set(j, 0, b.get(j, 0) - r.get(j, 0));
+            rr.set(j, 0, r.get(j, 0));
+        }
+
+        let bnrm = b.norm(T::two());
+        self.solve_preconditioner(a, &r, &mut z);
+
+        let mut p = M::zeros(n, 1);
+        let mut pp = M::zeros(n, 1);
+        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.get(j, 0) * rr.get(j, 0);
+            }
+            if iter == 1 {
+                for j in 0..n {
+                    p.set(j, 0, z.get(j, 0));
+                    pp.set(j, 0, zz.get(j, 0));
+                }
+            } else {
+                let bk = bknum / bkden;
+                for j in 0..n {
+                    p.set(j, 0, bk * p.get(j, 0) + z.get(j, 0));
+                    pp.set(j, 0, bk * pp.get(j, 0) + zz.get(j, 0));
+                }
+            }
+            bkden = bknum;
+            self.mat_vec_mul(a, &p, &mut z);
+            let mut akden = T::zero();
+            for j in 0..n {
+                akden += z.get(j, 0) * pp.get(j, 0);
+            }
+            let ak = bknum / akden;
+            self.mat_t_vec_mul(a, &pp, &mut zz);
+            for j in 0..n {
+                x.set(j, 0, x.get(j, 0) + ak * p.get(j, 0));
+                r.set(j, 0, r.get(j, 0) - ak * z.get(j, 0));
+                rr.set(j, 0, rr.get(j, 0) - ak * zz.get(j, 0));
+            }
+            self.solve_preconditioner(a, &r, &mut z);
+            err = r.norm(T::two()) / bnrm;
+
+            if err <= tol {
+                break;
+            }
+        }
+
+        Ok(err)
+    }
+
+    fn solve_preconditioner(&self, a: &M, b: &M, x: &mut M) {
+        let diag = Self::diag(a);
+        let n = diag.len();
+
+        for i in 0..n {
+            if diag[i] != T::zero() {
+                x.set(i, 0, b.get(i, 0) / diag[i]);
+            } else {
+                x.set(i, 0, b.get(i, 0));
+            }
+        }
+    }
+
+    // y = Ax
+    fn mat_vec_mul(&self, a: &M, x: &M, y: &mut M) {
+        y.copy_from(&a.matmul(x));
+    }
+
+    // y = Atx
+    fn mat_t_vec_mul(&self, a: &M, x: &M, y: &mut M) {
+        y.copy_from(&a.ab(true, x, false));
+    }
+
+    fn diag(a: &M) -> 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::naive::dense_matrix::*;
+
+    pub struct BGSolver {}
+
+    impl<T: RealNumber, M: Matrix<T>> BiconjugateGradientSolver<T, M> for BGSolver {}
+
+    #[test]
+    fn bg_solver() {
+        let a = DenseMatrix::from_2d_array(&[&[25., 15., -5.], &[15., 18., 0.], &[-5., 0., 11.]]);
+        let b = DenseMatrix::from_2d_array(&[&[40., 51., 28.]]);
+        let expected = DenseMatrix::from_2d_array(&[&[1.0, 2.0, 3.0]]);
+
+        let mut x = DenseMatrix::zeros(3, 1);
+
+        let solver = BGSolver {};
+
+        let err: f64 = solver
+            .solve_mut(&a, &b.transpose(), &mut x, 1e-6, 6)
+            .unwrap();
+
+        assert!(x.transpose().approximate_eq(&expected, 1e-4));
+        assert!((err - 0.0).abs() < 1e-4);
+    }
+}