impl Solve_Tridiagonal for Tridiagonal

Author: doraneko94

src/tridiagonal.rs

@@ -189,6 +189,57 @@ where
     }
 }
 
+impl<A> SolveTriDiagonal<A, Ix2> for TriDiagonal<A>
+where
+    A: Scalar + Lapack,
+{
+    fn solve_tridiagonal<Sb: Data<Elem = A>>(
+        &self,
+        b: &ArrayBase<Sb, Ix2>,
+    ) -> Result<Array<A, Ix2>> {
+        let mut b = replicate(b);
+        self.solve_tridiagonal_inplace(&mut b)?;
+        Ok(b)
+    }
+    fn solve_tridiagonal_into<Sb: DataMut<Elem = A>>(
+        &self,
+        mut b: ArrayBase<Sb, Ix2>,
+    ) -> Result<ArrayBase<Sb, Ix2>> {
+        self.solve_tridiagonal_inplace(&mut b)?;
+        Ok(b)
+    }
+    fn solve_t_tridiagonal<Sb: Data<Elem = A>>(
+        &self,
+        b: &ArrayBase<Sb, Ix2>,
+    ) -> Result<Array<A, Ix2>> {
+        let mut b = replicate(b);
+        self.solve_t_tridiagonal_inplace(&mut b)?;
+        Ok(b)
+    }
+    fn solve_t_tridiagonal_into<Sb: DataMut<Elem = A>>(
+        &self,
+        mut b: ArrayBase<Sb, Ix2>,
+    ) -> Result<ArrayBase<Sb, Ix2>> {
+        self.solve_t_tridiagonal_inplace(&mut b)?;
+        Ok(b)
+    }
+    fn solve_h_tridiagonal<Sb: Data<Elem = A>>(
+        &self,
+        b: &ArrayBase<Sb, Ix2>,
+    ) -> Result<Array<A, Ix2>> {
+        let mut b = replicate(b);
+        self.solve_h_tridiagonal_inplace(&mut b)?;
+        Ok(b)
+    }
+    fn solve_h_tridiagonal_into<Sb: DataMut<Elem = A>>(
+        &self,
+        mut b: ArrayBase<Sb, Ix2>,
+    ) -> Result<ArrayBase<Sb, Ix2>> {
+        self.solve_h_tridiagonal_inplace(&mut b)?;
+        Ok(b)
+    }
+}
+
 impl<A, S> SolveTriDiagonal<A, Ix2> for ArrayBase<S, Ix2>
 where
     A: Scalar + Lapack,
@@ -298,6 +349,42 @@ where
     }
 }
 
+impl<A> SolveTriDiagonalInplace<A, Ix2> for TriDiagonal<A>
+where
+    A: Scalar + Lapack,
+{
+    fn solve_tridiagonal_inplace<'a, Sb>(
+        &self,
+        rhs: &'a mut ArrayBase<Sb, Ix2>,
+    ) -> Result<&'a mut ArrayBase<Sb, Ix2>>
+    where
+        Sb: DataMut<Elem = A>,
+    {
+        let f = self.factorize_tridiagonal()?;
+        f.solve_tridiagonal_inplace(rhs)
+    }
+    fn solve_t_tridiagonal_inplace<'a, Sb>(
+        &self,
+        rhs: &'a mut ArrayBase<Sb, Ix2>,
+    ) -> Result<&'a mut ArrayBase<Sb, Ix2>>
+    where
+        Sb: DataMut<Elem = A>,
+    {
+        let f = self.factorize_tridiagonal()?;
+        f.solve_t_tridiagonal_inplace(rhs)
+    }
+    fn solve_h_tridiagonal_inplace<'a, Sb>(
+        &self,
+        rhs: &'a mut ArrayBase<Sb, Ix2>,
+    ) -> Result<&'a mut ArrayBase<Sb, Ix2>>
+    where
+        Sb: DataMut<Elem = A>,
+    {
+        let f = self.factorize_tridiagonal()?;
+        f.solve_h_tridiagonal_inplace(rhs)
+    }
+}
+
 impl<A, S> SolveTriDiagonalInplace<A, Ix2> for ArrayBase<S, Ix2>
 where
     A: Scalar + Lapack,
@@ -383,6 +470,60 @@ where
     }
 }
 
+impl<A> SolveTriDiagonal<A, Ix1> for TriDiagonal<A>
+where
+    A: Scalar + Lapack,
+{
+    fn solve_tridiagonal<Sb: Data<Elem = A>>(
+        &self,
+        b: &ArrayBase<Sb, Ix1>,
+    ) -> Result<Array<A, Ix1>> {
+        let b = b.to_owned();
+        self.solve_tridiagonal_into(b)
+    }
+    fn solve_tridiagonal_into<Sb: DataMut<Elem = A>>(
+        &self,
+        b: ArrayBase<Sb, Ix1>,
+    ) -> Result<ArrayBase<Sb, Ix1>> {
+        let b = into_col(b);
+        let f = self.factorize_tridiagonal()?;
+        let b = f.solve_tridiagonal_into(b)?;
+        Ok(flatten(b))
+    }
+    fn solve_t_tridiagonal<Sb: Data<Elem = A>>(
+        &self,
+        b: &ArrayBase<Sb, Ix1>,
+    ) -> Result<Array<A, Ix1>> {
+        let b = b.to_owned();
+        self.solve_t_tridiagonal_into(b)
+    }
+    fn solve_t_tridiagonal_into<Sb: DataMut<Elem = A>>(
+        &self,
+        b: ArrayBase<Sb, Ix1>,
+    ) -> Result<ArrayBase<Sb, Ix1>> {
+        let b = into_col(b);
+        let f = self.factorize_tridiagonal()?;
+        let b = f.solve_t_tridiagonal_into(b)?;
+        Ok(flatten(b))
+    }
+    fn solve_h_tridiagonal<Sb: Data<Elem = A>>(
+        &self,
+        b: &ArrayBase<Sb, Ix1>,
+    ) -> Result<Array<A, Ix1>> {
+        let b = b.to_owned();
+        self.solve_h_tridiagonal_into(b)
+    }
+    fn solve_h_tridiagonal_into<Sb: DataMut<Elem = A>>(
+        &self,
+        b: ArrayBase<Sb, Ix1>,
+    ) -> Result<ArrayBase<Sb, Ix1>> {
+        let b = into_col(b);
+        let f = self.factorize_tridiagonal()?;
+        let b = f.solve_h_tridiagonal_into(b)?;
+        Ok(flatten(b))
+    }
+}
+
 impl<A, S> SolveTriDiagonal<A, Ix1> for ArrayBase<S, Ix2>
 where
     A: Scalar + Lapack,

tests/tridiagonal.rs

@@ -125,6 +125,21 @@ fn solve_tridiagonal_random_t() {
     assert_close_l2!(&y1, &y2, 1e-7);
 }
 
+#[test]
+fn to_tridiagonal_solve_random() {
+    let mut a: Array2<f64> = random((3, 3));
+    a[[0, 2]] = 0.0;
+    a[[2, 0]] = 0.0;
+    let tridiag = a.to_tridiagonal().unwrap();
+    let x: Array1<f64> = random(3);
+    let b1 = a.dot(&x);
+    let b2 = b1.clone();
+    let y1 = tridiag.solve_tridiagonal_into(b1).unwrap();
+    let y2 = a.solve_into(b2).unwrap();
+    assert_close_l2!(&x, &y1, 1e-7);
+    assert_close_l2!(&y1, &y2, 1e-7);
+}
+
 #[test]
 fn det_tridiagonal_f64() {
     let a: Array2<f64> = arr2(&[[10.0, -9.0, 0.0], [7.0, -12.0, 11.0], [0.0, 10.0, 3.0]]);