Revise interface of solve()

examples/solve.rs

@@ -5,17 +5,26 @@ extern crate ndarray_linalg;
 use ndarray::*;
 use ndarray_linalg::*;
 
+// Solve `Ax=b`
+fn solve() -> Result<(), error::LinalgError> {
+    let a: Array2<f64> = random((3, 3));
+    let b: Array1<f64> = random(3);
+    let _x = a.solve(&b)?;
+    Ok(())
+}
+
 // Solve `Ax=b` for many b with fixed A
 fn factorize() -> Result<(), error::LinalgError> {
     let a: Array2<f64> = random((3, 3));
     let f = a.factorize_into()?; // LU factorize A (A is consumed)
     for _ in 0..10 {
         let b: Array1<f64> = random(3);
-        let x = f.solve(Transpose::No, b)?; // solve Ax=b using factorized L, U
+        let _x = f.solve_into(b)?; // solve Ax=b using factorized L, U
     }
     Ok(())
 }
 
 fn main() {
+    solve().unwrap();
     factorize().unwrap();
 }

src/solve.rs

@@ -9,24 +9,89 @@ use super::types::*;
 
 pub use lapack_traits::{Pivot, Transpose};
 
+pub trait Solve<A: Scalar> {
+    fn solve<S: Data<Elem = A>>(&self, a: &ArrayBase<S, Ix1>) -> Result<Array1<A>> {
+        let mut a = replicate(a);
+        self.solve_mut(&mut a)?;
+        Ok(a)
+    }
+    fn solve_into<S: DataMut<Elem = A>>(&self, mut a: ArrayBase<S, Ix1>) -> Result<ArrayBase<S, Ix1>> {
+        self.solve_mut(&mut a)?;
+        Ok(a)
+    }
+    fn solve_mut<'a, S: DataMut<Elem = A>>(&self, &'a mut ArrayBase<S, Ix1>) -> Result<&'a mut ArrayBase<S, Ix1>>;
+
+    fn solve_t<S: Data<Elem = A>>(&self, a: &ArrayBase<S, Ix1>) -> Result<Array1<A>> {
+        let mut a = replicate(a);
+        self.solve_t_mut(&mut a)?;
+        Ok(a)
+    }
+    fn solve_t_into<S: DataMut<Elem = A>>(&self, mut a: ArrayBase<S, Ix1>) -> Result<ArrayBase<S, Ix1>> {
+        self.solve_t_mut(&mut a)?;
+        Ok(a)
+    }
+    fn solve_t_mut<'a, S: DataMut<Elem = A>>(&self, &'a mut ArrayBase<S, Ix1>) -> Result<&'a mut ArrayBase<S, Ix1>>;
+
+    fn solve_h<S: Data<Elem = A>>(&self, a: &ArrayBase<S, Ix1>) -> Result<Array1<A>> {
+        let mut a = replicate(a);
+        self.solve_h_mut(&mut a)?;
+        Ok(a)
+    }
+    fn solve_h_into<S: DataMut<Elem = A>>(&self, mut a: ArrayBase<S, Ix1>) -> Result<ArrayBase<S, Ix1>> {
+        self.solve_h_mut(&mut a)?;
+        Ok(a)
+    }
+    fn solve_h_mut<'a, S: DataMut<Elem = A>>(&self, &'a mut ArrayBase<S, Ix1>) -> Result<&'a mut ArrayBase<S, Ix1>>;
+}
+
 pub struct Factorized<S: Data> {
     pub a: ArrayBase<S, Ix2>,
     pub ipiv: Pivot,
 }
 
-impl<A, S> Factorized<S>
+impl<A, S> Solve<A> for Factorized<S>
 where
     A: Scalar,
     S: Data<Elem = A>,
 {
-    pub fn solve<Sb>(&self, t: Transpose, mut rhs: ArrayBase<Sb, Ix1>) -> Result<ArrayBase<Sb, Ix1>>
+    fn solve_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
+    where
+        Sb: DataMut<Elem = A>,
+    {
+        unsafe {
+            A::solve(
+                self.a.square_layout()?,
+                Transpose::No,
+                self.a.as_allocated()?,
+                &self.ipiv,
+                rhs.as_slice_mut().unwrap(),
+            )?
+        };
+        Ok(rhs)
+    }
+    fn solve_t_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
+    where
+        Sb: DataMut<Elem = A>,
+    {
+        unsafe {
+            A::solve(
+                self.a.square_layout()?,
+                Transpose::Transpose,
+                self.a.as_allocated()?,
+                &self.ipiv,
+                rhs.as_slice_mut().unwrap(),
+            )?
+        };
+        Ok(rhs)
+    }
+    fn solve_h_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
     where
         Sb: DataMut<Elem = A>,
     {
         unsafe {
             A::solve(
                 self.a.square_layout()?,
-                t,
+                Transpose::Hermite,
                 self.a.as_allocated()?,
                 &self.ipiv,
                 rhs.as_slice_mut().unwrap(),
@@ -36,6 +101,34 @@ where
     }
 }
 
+impl<A, S> Solve<A> for ArrayBase<S, Ix2>
+where
+    A: Scalar,
+    S: Data<Elem = A>,
+{
+    fn solve_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
+    where
+        Sb: DataMut<Elem = A>,
+    {
+        let f = self.factorize()?;
+        f.solve_mut(rhs)
+    }
+    fn solve_t_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
+    where
+        Sb: DataMut<Elem = A>,
+    {
+        let f = self.factorize()?;
+        f.solve_t_mut(rhs)
+    }
+    fn solve_h_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
+    where
+        Sb: DataMut<Elem = A>,
+    {
+        let f = self.factorize()?;
+        f.solve_h_mut(rhs)
+    }
+}
+
 impl<A, S> Factorized<S>
 where
     A: Scalar,