Re-implement SolveTriangular

Author: termoshtt

src/error.rs

@@ -2,7 +2,7 @@
 
 use std::error;
 use std::fmt;
-use ndarray::Ixs;
+use ndarray::{Ixs, ShapeError};
 
 #[derive(Debug)]
 pub struct LapackError {
@@ -68,6 +68,7 @@ pub enum LinalgError {
     NotSquare(NotSquareError),
     Lapack(LapackError),
     Stride(StrideError),
+    Shape(ShapeError),
 }
 
 impl fmt::Display for LinalgError {
@@ -76,6 +77,7 @@ impl fmt::Display for LinalgError {
             LinalgError::NotSquare(ref err) => err.fmt(f),
             LinalgError::Lapack(ref err) => err.fmt(f),
             LinalgError::Stride(ref err) => err.fmt(f),
+            LinalgError::Shape(ref err) => err.fmt(f),
         }
     }
 }
@@ -86,6 +88,7 @@ impl error::Error for LinalgError {
             LinalgError::NotSquare(ref err) => err.description(),
             LinalgError::Lapack(ref err) => err.description(),
             LinalgError::Stride(ref err) => err.description(),
+            LinalgError::Shape(ref err) => err.description(),
         }
     }
 }
@@ -107,3 +110,8 @@ impl From<StrideError> for LinalgError {
         LinalgError::Stride(err)
     }
 }
+impl From<ShapeError> for LinalgError {
+    fn from(err: ShapeError) -> LinalgError {
+        LinalgError::Shape(err)
+    }
+}

src/solve.rs

@@ -18,13 +18,13 @@ pub trait ImplSolve: Sized {
              b: Vec<Self>)
              -> Result<Vec<Self>, LapackError>;
     /// solve triangular linear problem
-    fn solve_triangle(layout: Layout,
-                      uplo: u8,
-                      size: usize,
-                      a: &[Self],
-                      b: Vec<Self>,
-                      nrhs: i32)
-                      -> Result<Vec<Self>, LapackError>;
+    fn solve_triangle<'a, 'b>(layout: Layout,
+                              uplo: u8,
+                              size: usize,
+                              a: &'a [Self],
+                              b: &'b mut [Self],
+                              nrhs: i32)
+                              -> Result<&'b mut [Self], LapackError>;
 }
 
 macro_rules! impl_solve {
@@ -66,7 +66,7 @@ impl ImplSolve for $scalar {
             Err(From::from(info))
         }
     }
-    fn solve_triangle(layout: Layout, uplo: u8, size: usize, a: &[Self], mut b: Vec<Self>, nrhs: i32) -> Result<Vec<Self>, LapackError> {
+    fn solve_triangle<'a, 'b>(layout: Layout, uplo: u8, size: usize, a: &'a [Self], mut b: &'b mut [Self], nrhs: i32) -> Result<&'b mut [Self], LapackError> {
         let n = size as i32;
         let lda = n;
         let ldb = match layout {

src/triangular.rs

@@ -1,10 +1,10 @@
 
-use ndarray::{Ix1, Ix2, Array, RcArray, NdFloat, ArrayBase, DataMut};
-
-use matrix::{Matrix, MFloat};
-use square::SquareMatrix;
-use error::LinalgError;
-use solve::ImplSolve;
+use ndarray::*;
+use super::matrix::{Matrix, MFloat};
+use super::square::SquareMatrix;
+use super::error::LinalgError;
+use super::solve::ImplSolve;
+use super::util::hstack;
 
 pub trait SolveTriangular<Rhs>: Matrix + SquareMatrix {
     type Output;
@@ -14,34 +14,70 @@ pub trait SolveTriangular<Rhs>: Matrix + SquareMatrix {
     fn solve_lower(&self, Rhs) -> Result<Self::Output, LinalgError>;
 }
 
-impl<A: MFloat> SolveTriangular<Array<A, Ix1>> for Array<A, Ix2> {
-    type Output = Array<A, Ix1>;
-    fn solve_upper(&self, b: Array<A, Ix1>) -> Result<Self::Output, LinalgError> {
+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
+{
+    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();
-        let x = ImplSolve::solve_triangle(layout, 'U' as u8, n, a, b.into_raw_vec(), 1)?;
-        Ok(Array::from_vec(x))
+        {
+            let b_ = b.as_slice_memory_order_mut().unwrap();
+            ImplSolve::solve_triangle(layout, 'U' as u8, n, a, b_, 1)?;
+        }
+        Ok(b)
     }
-    fn solve_lower(&self, b: Array<A, Ix1>) -> Result<Self::Output, LinalgError> {
+    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();
-        let x = ImplSolve::solve_triangle(layout, 'L' as u8, n, a, b.into_raw_vec(), 1)?;
-        Ok(Array::from_vec(x))
+        {
+            let b_ = b.as_slice_memory_order_mut().unwrap();
+            ImplSolve::solve_triangle(layout, 'L' as u8, n, a, b_, 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())
     }
 }
 
-impl<A: MFloat> SolveTriangular<RcArray<A, Ix1>> for RcArray<A, Ix2> {
-    type Output = RcArray<A, Ix1>;
-    fn solve_upper(&self, b: RcArray<A, Ix1>) -> Result<Self::Output, LinalgError> {
+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.into_owned())?;
+        let x = self.to_owned().solve_upper(&b)?;
         Ok(x.into_shared())
     }
-    fn solve_lower(&self, b: RcArray<A, Ix1>) -> Result<Self::Output, LinalgError> {
+    fn solve_lower(&self, b: RcArray<A, Ix2>) -> Result<Self::Output, LinalgError> {
         // XXX unnecessary clone
-        let x = self.to_owned().solve_lower(b.into_owned())?;
+        let x = self.to_owned().solve_lower(&b)?;
         Ok(x.into_shared())
     }
 }