impl AllocatedArrayMut for ArrayBase1 to use triangular

Author: termoshtt

src/impl2/mod.rs

@@ -13,11 +13,14 @@ pub use self::svd::*;
 pub use self::solve::*;
 pub use self::cholesky::*;
 pub use self::eigh::*;
+pub use self::triangular::*;
 
 use super::error::*;
 
-pub trait LapackScalar: OperatorNorm_ + QR_ + SVD_ + Solve_ + Cholesky_ + Eigh_ {}
-impl<A> LapackScalar for A where A: OperatorNorm_ + QR_ + SVD_ + Solve_ + Cholesky_ + Eigh_ {}
+pub trait LapackScalar
+    : OperatorNorm_ + QR_ + SVD_ + Solve_ + Cholesky_ + Eigh_ + Triangular_ {
+}
+impl<A> LapackScalar for A where A: OperatorNorm_ + QR_ + SVD_ + Solve_ + Cholesky_ + Eigh_ + Triangular_ {}
 
 pub fn into_result<T>(info: i32, val: T) -> Result<T> {
     if info == 0 {

src/impl2/triangular.rs

@@ -8,31 +8,25 @@ use super::{UPLO, Transpose, into_result};
 
 #[derive(Debug, Clone, Copy)]
 #[repr(u8)]
-pub enum TriangularDiag {
+pub enum Diag {
     Unit = b'U',
     NonUnit = b'N',
 }
 
 pub trait Triangular_: Sized {
-    fn inv_triangular(l: Layout, UPLO, TriangularDiag, a: &mut [Self]) -> Result<()>;
-    fn solve_triangular(al: Layout, bl: Layout, UPLO, TriangularDiag, a: &[Self], b: &mut [Self]) -> Result<()>;
+    fn inv_triangular(l: Layout, UPLO, Diag, a: &mut [Self]) -> Result<()>;
+    fn solve_triangular(al: Layout, bl: Layout, UPLO, Diag, a: &[Self], b: &mut [Self]) -> Result<()>;
 }
 
 impl Triangular_ for f64 {
-    fn inv_triangular(l: Layout, uplo: UPLO, diag: TriangularDiag, a: &mut [Self]) -> Result<()> {
+    fn inv_triangular(l: Layout, uplo: UPLO, diag: Diag, a: &mut [Self]) -> Result<()> {
         let (n, _) = l.size();
         let lda = l.lda();
         let info = c::dtrtri(l.lapacke_layout(), uplo as u8, diag as u8, n, a, lda);
         into_result(info, ())
     }
 
-    fn solve_triangular(al: Layout,
-                        bl: Layout,
-                        uplo: UPLO,
-                        diag: TriangularDiag,
-                        a: &[Self],
-                        mut b: &mut [Self])
-                        -> Result<()> {
+    fn solve_triangular(al: Layout, bl: Layout, uplo: UPLO, diag: Diag, a: &[Self], mut b: &mut [Self]) -> Result<()> {
         let (n, _) = al.size();
         let lda = al.lda();
         let nrhs = bl.len();

src/layout.rs

@@ -53,20 +53,20 @@ impl Layout {
 }
 
 pub trait AllocatedArray {
-    type Scalar;
+    type Elem;
     fn layout(&self) -> Result<Layout>;
     fn square_layout(&self) -> Result<Layout>;
-    fn as_allocated(&self) -> Result<&[Self::Scalar]>;
+    fn as_allocated(&self) -> Result<&[Self::Elem]>;
 }
 
 pub trait AllocatedArrayMut: AllocatedArray {
-    fn as_allocated_mut(&mut self) -> Result<&mut [Self::Scalar]>;
+    fn as_allocated_mut(&mut self) -> Result<&mut [Self::Elem]>;
 }
 
 impl<A, S> AllocatedArray for ArrayBase<S, Ix2>
     where S: Data<Elem = A>
 {
-    type Scalar = A;
+    type Elem = A;
 
     fn layout(&self) -> Result<Layout> {
         let strides = self.strides();
@@ -91,20 +91,45 @@ impl<A, S> AllocatedArray for ArrayBase<S, Ix2>
     }
 
     fn as_allocated(&self) -> Result<&[A]> {
-        let slice = self.as_slice_memory_order().ok_or(MemoryContError::new())?;
-        Ok(slice)
+        Ok(self.as_slice_memory_order().ok_or(MemoryContError::new())?)
     }
 }
 
 impl<A, S> AllocatedArrayMut for ArrayBase<S, Ix2>
     where S: DataMut<Elem = A>
 {
     fn as_allocated_mut(&mut self) -> Result<&mut [A]> {
-        let slice = self.as_slice_memory_order_mut().ok_or(MemoryContError::new())?;
-        Ok(slice)
+        Ok(self.as_slice_memory_order_mut().ok_or(MemoryContError::new())?)
     }
 }
 
+impl<A, S> AllocatedArray for ArrayBase<S, Ix1>
+    where S: Data<Elem = A>
+{
+    type Elem = A;
+
+    fn layout(&self) -> Result<Layout> {
+        Ok(Layout::F((self.len() as i32, 1)))
+    }
+
+    fn square_layout(&self) -> Result<Layout> {
+        Err(NotSquareError::new(self.len() as i32, 1).into())
+    }
+
+    fn as_allocated(&self) -> Result<&[A]> {
+        Ok(self.as_slice_memory_order().ok_or(MemoryContError::new())?)
+    }
+}
+
+impl<A, S> AllocatedArrayMut for ArrayBase<S, Ix1>
+    where S: DataMut<Elem = A>
+{
+    fn as_allocated_mut(&mut self) -> Result<&mut [A]> {
+        Ok(self.as_slice_memory_order_mut().ok_or(MemoryContError::new())?)
+    }
+}
+
+
 pub fn reconstruct<A, S>(l: Layout, a: Vec<A>) -> Result<ArrayBase<S, Ix2>>
     where S: DataOwned<Elem = A>
 {

src/triangular.rs

@@ -2,114 +2,43 @@
 
 use ndarray::*;
 use num_traits::Zero;
-use super::impl2::UPLO;
 
-use super::matrix::{Matrix, MFloat};
-use super::square::SquareMatrix;
-use super::error::LinalgError;
-use super::util::hstack;
-use super::impls::solve::ImplSolve;
+use super::layout::*;
+use super::error::*;
+use super::impl2::*;
 
-pub trait SolveTriangular<Rhs>: Matrix + SquareMatrix {
+/// solve a triangular system with upper triangular matrix
+pub trait SolveTriangular<Rhs> {
     type Output;
-    /// solve a triangular system with upper triangular matrix
-    fn solve_upper(&self, Rhs) -> Result<Self::Output, LinalgError>;
-    /// solve a triangular system with lower triangular matrix
-    fn solve_lower(&self, Rhs) -> Result<Self::Output, LinalgError>;
+    fn solve_triangular(&self, UPLO, Diag, Rhs) -> Result<Self::Output>;
 }
 
-impl<A, S1, S2> SolveTriangular<ArrayBase<S2, Ix1>> for ArrayBase<S1, Ix2>
-    where A: MFloat,
-          S1: Data<Elem = A>,
-          S2: DataMut<Elem = A>,
-          ArrayBase<S1, Ix2>: Matrix + SquareMatrix
+impl<A, S, V> SolveTriangular<V> for ArrayBase<S, Ix2>
+    where A: LapackScalar,
+          S: Data<Elem = A>,
+          V: AllocatedArrayMut<Elem = A>
 {
-    type Output = ArrayBase<S2, Ix1>;
-    fn solve_upper(&self, mut b: ArrayBase<S2, Ix1>) -> Result<Self::Output, LinalgError> {
-        let n = self.square_size()?;
-        let layout = self.layout()?;
-        let a = self.as_slice_memory_order().unwrap();
-        ImplSolve::solve_triangle(layout,
-                                  'U' as u8,
-                                  n,
-                                  a,
-                                  b.as_slice_memory_order_mut().unwrap(),
-                                  1)?;
-        Ok(b)
-    }
-    fn solve_lower(&self, mut b: ArrayBase<S2, Ix1>) -> Result<Self::Output, LinalgError> {
-        let n = self.square_size()?;
-        let layout = self.layout()?;
-        let a = self.as_slice_memory_order().unwrap();
-        ImplSolve::solve_triangle(layout,
-                                  'L' as u8,
-                                  n,
-                                  a,
-                                  b.as_slice_memory_order_mut().unwrap(),
-                                  1)?;
-        Ok(b)
-    }
-}
-
-impl<'a, S1, S2, A> SolveTriangular<&'a ArrayBase<S2, Ix2>> for ArrayBase<S1, Ix2>
-    where A: MFloat,
-          S1: Data<Elem = A>,
-          S2: Data<Elem = A>,
-          ArrayBase<S1, Ix2>: Matrix + SquareMatrix
-{
-    type Output = Array<A, Ix2>;
-    fn solve_upper(&self, bs: &ArrayBase<S2, Ix2>) -> Result<Self::Output, LinalgError> {
-        let mut xs = Vec::new();
-        for b in bs.axis_iter(Axis(1)) {
-            let x = self.solve_upper(b.to_owned())?;
-            xs.push(x);
-        }
-        hstack(&xs).map_err(|e| e.into())
-    }
-    fn solve_lower(&self, bs: &ArrayBase<S2, Ix2>) -> Result<Self::Output, LinalgError> {
-        let mut xs = Vec::new();
-        for b in bs.axis_iter(Axis(1)) {
-            let x = self.solve_lower(b.to_owned())?;
-            xs.push(x);
-        }
-        hstack(&xs).map_err(|e| e.into())
-    }
-}
+    type Output = V;
 
-impl<A: MFloat> SolveTriangular<RcArray<A, Ix2>> for RcArray<A, Ix2> {
-    type Output = RcArray<A, Ix2>;
-    fn solve_upper(&self, b: RcArray<A, Ix2>) -> Result<Self::Output, LinalgError> {
-        // XXX unnecessary clone
-        let x = self.to_owned().solve_upper(&b)?;
-        Ok(x.into_shared())
-    }
-    fn solve_lower(&self, b: RcArray<A, Ix2>) -> Result<Self::Output, LinalgError> {
-        // XXX unnecessary clone
-        let x = self.to_owned().solve_lower(&b)?;
-        Ok(x.into_shared())
+    fn solve_triangular(&self, uplo: UPLO, diag: Diag, mut b: V) -> Result<Self::Output> {
+        let la = self.layout()?;
+        let lb = b.layout()?;
+        let a_ = self.as_allocated()?;
+        A::solve_triangular(la, lb, uplo, diag, a_, b.as_allocated_mut()?)?;
+        Ok(b)
     }
 }
 
-pub fn drop_upper<A: Zero, S>(mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
+pub fn drop_upper<A: Zero, S>(a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
     where S: DataMut<Elem = A>
 {
-    for ((i, j), val) in a.indexed_iter_mut() {
-        if i < j {
-            *val = A::zero();
-        }
-    }
-    a
+    a.into_triangular(UPLO::Lower)
 }
 
-pub fn drop_lower<A: Zero, S>(mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
+pub fn drop_lower<A: Zero, S>(a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
     where S: DataMut<Elem = A>
 {
-    for ((i, j), val) in a.indexed_iter_mut() {
-        if i > j {
-            *val = A::zero();
-        }
-    }
-    a
+    a.into_triangular(UPLO::Upper)
 }
 
 pub trait IntoTriangular<T> {