Add IntoTriangular

Author: termoshtt

src/cholesky.rs

@@ -1,8 +1,10 @@
 
 use ndarray::*;
+use num_traits::Zero;
 
 use super::error::*;
 use super::layout::*;
+use super::triangular::IntoTriangular;
 
 use impl2::LapackScalar;
 pub use impl2::UPLO;
@@ -12,33 +14,33 @@ pub trait Cholesky<K> {
 }
 
 impl<A, S> Cholesky<ArrayBase<S, Ix2>> for ArrayBase<S, Ix2>
-    where A: LapackScalar,
+    where A: LapackScalar + Zero,
           S: DataMut<Elem = A>
 {
     fn cholesky(mut self, uplo: UPLO) -> Result<ArrayBase<S, Ix2>> {
         A::cholesky(self.square_layout()?, uplo, self.as_allocated_mut()?)?;
-        Ok(self)
+        Ok(self.into_triangular(uplo))
     }
 }
 
 impl<'a, A, S> Cholesky<&'a mut ArrayBase<S, Ix2>> for &'a mut ArrayBase<S, Ix2>
-    where A: LapackScalar,
+    where A: LapackScalar + Zero,
           S: DataMut<Elem = A>
 {
     fn cholesky(mut self, uplo: UPLO) -> Result<&'a mut ArrayBase<S, Ix2>> {
         A::cholesky(self.square_layout()?, uplo, self.as_allocated_mut()?)?;
-        Ok(self)
+        Ok(self.into_triangular(uplo))
     }
 }
 
 impl<'a, A, Si, So> Cholesky<ArrayBase<So, Ix2>> for &'a ArrayBase<Si, Ix2>
-    where A: LapackScalar + Copy,
+    where A: LapackScalar + Copy + Zero,
           Si: Data<Elem = A>,
           So: DataMut<Elem = A> + DataOwned
 {
     fn cholesky(self, uplo: UPLO) -> Result<ArrayBase<So, Ix2>> {
         let mut a = replicate(self);
         A::cholesky(a.square_layout()?, uplo, a.as_allocated_mut()?)?;
-        Ok(a)
+        Ok(a.into_triangular(uplo))
     }
 }

src/triangular.rs

@@ -1,6 +1,9 @@
 //! Define methods for triangular matrices
 
 use ndarray::*;
+use num_traits::Zero;
+use super::impl2::UPLO;
+
 use super::matrix::{Matrix, MFloat};
 use super::square::SquareMatrix;
 use super::error::LinalgError;
@@ -87,7 +90,7 @@ impl<A: MFloat> SolveTriangular<RcArray<A, Ix2>> for RcArray<A, Ix2> {
     }
 }
 
-pub fn drop_upper<A: NdFloat, S>(mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
+pub fn drop_upper<A: Zero, S>(mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
     where S: DataMut<Elem = A>
 {
     for ((i, j), val) in a.indexed_iter_mut() {
@@ -98,7 +101,7 @@ pub fn drop_upper<A: NdFloat, S>(mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
     a
 }
 
-pub fn drop_lower<A: NdFloat, S>(mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
+pub fn drop_lower<A: Zero, S>(mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
     where S: DataMut<Elem = A>
 {
     for ((i, j), val) in a.indexed_iter_mut() {
@@ -108,3 +111,42 @@ pub fn drop_lower<A: NdFloat, S>(mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
     }
     a
 }
+
+pub trait IntoTriangular<T> {
+    fn into_triangular(self, UPLO) -> T;
+}
+
+impl<'a, A, S> IntoTriangular<&'a mut ArrayBase<S, Ix2>> for &'a mut ArrayBase<S, Ix2>
+    where A: Zero,
+          S: DataMut<Elem = A>
+{
+    fn into_triangular(self, uplo: UPLO) -> &'a mut ArrayBase<S, Ix2> {
+        match uplo {
+            UPLO::Upper => {
+                for ((i, j), val) in self.indexed_iter_mut() {
+                    if i > j {
+                        *val = A::zero();
+                    }
+                }
+            }
+            UPLO::Lower => {
+                for ((i, j), val) in self.indexed_iter_mut() {
+                    if i < j {
+                        *val = A::zero();
+                    }
+                }
+            }
+        }
+        self
+    }
+}
+
+impl<A, S> IntoTriangular<ArrayBase<S, Ix2>> for ArrayBase<S, Ix2>
+    where A: Zero,
+          S: DataMut<Elem = A>
+{
+    fn into_triangular(mut self, uplo: UPLO) -> ArrayBase<S, Ix2> {
+        (&mut self).into_triangular(uplo);
+        self
+    }
+}