Replace operator::Operator* by operator::LinearOperator

Author: termoshtt

src/diagonal.rs

@@ -2,8 +2,8 @@
 
 use ndarray::*;
 
-use super::convert::*;
 use super::operator::*;
+use super::types::*;
 
 /// Vector as a Diagonal matrix
 pub struct Diagonal<S: Data> {
@@ -30,81 +30,19 @@ impl<A, S: Data<Elem = A>> AsDiagonal<A> for ArrayBase<S, Ix1> {
     }
 }
 
-impl<A, S, Sr> OperatorInplace<Sr, Ix1> for Diagonal<S>
+impl<A, Sa> LinearOperator for Diagonal<Sa>
 where
-    A: LinalgScalar,
-    S: Data<Elem = A>,
-    Sr: DataMut<Elem = A>,
+    A: Scalar,
+    Sa: Data<Elem = A>,
 {
-    fn op_inplace<'a>(&self, a: &'a mut ArrayBase<Sr, Ix1>) -> &'a mut ArrayBase<Sr, Ix1> {
+    type Elem = A;
+
+    fn apply_mut<S>(&self, a: &mut ArrayBase<S, Ix1>)
+    where
+        S: DataMut<Elem = A>,
+    {
         for (val, d) in a.iter_mut().zip(self.diag.iter()) {
             *val = *val * *d;
         }
-        a
-    }
-}
-
-impl<A, S, Sr> Operator<A, Sr, Ix1> for Diagonal<S>
-where
-    A: LinalgScalar,
-    S: Data<Elem = A>,
-    Sr: Data<Elem = A>,
-{
-    fn op(&self, a: &ArrayBase<Sr, Ix1>) -> Array1<A> {
-        let mut a = replicate(a);
-        self.op_inplace(&mut a);
-        a
-    }
-}
-
-impl<A, S, Sr> OperatorInto<Sr, Ix1> for Diagonal<S>
-where
-    A: LinalgScalar,
-    S: Data<Elem = A>,
-    Sr: DataOwned<Elem = A> + DataMut,
-{
-    fn op_into(&self, mut a: ArrayBase<Sr, Ix1>) -> ArrayBase<Sr, Ix1> {
-        self.op_inplace(&mut a);
-        a
-    }
-}
-
-impl<A, S, Sr> OperatorInplace<Sr, Ix2> for Diagonal<S>
-where
-    A: LinalgScalar,
-    S: Data<Elem = A>,
-    Sr: DataMut<Elem = A>,
-{
-    fn op_inplace<'a>(&self, a: &'a mut ArrayBase<Sr, Ix2>) -> &'a mut ArrayBase<Sr, Ix2> {
-        let d = &self.diag;
-        for ((i, _), val) in a.indexed_iter_mut() {
-            *val = *val * d[i];
-        }
-        a
-    }
-}
-
-impl<A, S, Sr> Operator<A, Sr, Ix2> for Diagonal<S>
-where
-    A: LinalgScalar,
-    S: Data<Elem = A>,
-    Sr: Data<Elem = A>,
-{
-    fn op(&self, a: &ArrayBase<Sr, Ix2>) -> Array2<A> {
-        let mut a = replicate(a);
-        self.op_inplace(&mut a);
-        a
-    }
-}
-
-impl<A, S, Sr> OperatorInto<Sr, Ix2> for Diagonal<S>
-where
-    A: LinalgScalar,
-    S: Data<Elem = A>,
-    Sr: DataOwned<Elem = A> + DataMut,
-{
-    fn op_into(&self, mut a: ArrayBase<Sr, Ix2>) -> ArrayBase<Sr, Ix2> {
-        self.op_inplace(&mut a);
-        a
     }
 }

src/eigh.rs

@@ -5,7 +5,7 @@ use ndarray::*;
 use crate::diagonal::*;
 use crate::error::*;
 use crate::layout::*;
-use crate::operator::Operator;
+use crate::operator::LinearOperator;
 use crate::types::*;
 use crate::UPLO;
 
@@ -165,7 +165,7 @@ where
     fn ssqrt_into(self, uplo: UPLO) -> Result<Self::Output> {
         let (e, v) = self.eigh_into(uplo)?;
         let e_sqrt = Array1::from_iter(e.iter().map(|r| Scalar::from_real(r.sqrt())));
-        let ev = e_sqrt.into_diagonal().op(&v.t());
-        Ok(v.op(&ev))
+        let ev = e_sqrt.into_diagonal().apply2(&v.t());
+        Ok(v.apply2(&ev))
     }
 }

src/krylov/mod.rs

@@ -6,7 +6,6 @@ use ndarray::*;
 pub mod arnoldi;
 pub mod householder;
 pub mod mgs;
-pub mod operator;
 
 pub use arnoldi::{arnoldi_householder, arnoldi_mgs, Arnoldi};
 pub use householder::{householder, Householder};

src/krylov/operator.rs

@@ -1,105 +1,80 @@
-//! Linear Operator
+//! Linear operator algebra
 
+use crate::generate::hstack;
+use crate::types::*;
 use ndarray::*;
 
-use super::types::*;
+pub trait LinearOperator {
+    type Elem: Scalar;
 
-pub trait Operator<A, S, D>
-where
-    S: Data<Elem = A>,
-    D: Dimension,
-{
-    fn op(&self, a: &ArrayBase<S, D>) -> Array<A, D>;
-}
-
-pub trait OperatorInto<S, D>
-where
-    S: DataMut,
-    D: Dimension,
-{
-    fn op_into(&self, a: ArrayBase<S, D>) -> ArrayBase<S, D>;
-}
-
-pub trait OperatorInplace<S, D>
-where
-    S: DataMut,
-    D: Dimension,
-{
-    fn op_inplace<'a>(&self, a: &'a mut ArrayBase<S, D>) -> &'a mut ArrayBase<S, D>;
-}
+    /// Apply operator out-place
+    fn apply<S>(&self, a: &ArrayBase<S, Ix1>) -> Array1<S::Elem>
+    where
+        S: Data<Elem = Self::Elem>,
+    {
+        let mut a = a.to_owned();
+        self.apply_mut(&mut a);
+        a
+    }
 
-impl<T, A, S, D> Operator<A, S, D> for T
-where
-    A: Scalar + Lapack,
-    S: Data<Elem = A>,
-    D: Dimension,
-    T: linalg::Dot<ArrayBase<S, D>, Output = Array<A, D>>,
-{
-    fn op(&self, rhs: &ArrayBase<S, D>) -> Array<A, D> {
-        self.dot(rhs)
+    /// Apply operator in-place
+    fn apply_mut<S>(&self, a: &mut ArrayBase<S, Ix1>)
+    where
+        S: DataMut<Elem = Self::Elem>,
+    {
+        let b = self.apply(a);
+        azip!(mut a(a), b in { *a = b });
     }
-}
 
-pub trait OperatorMulti<A, S, D>
-where
-    S: Data<Elem = A>,
-    D: Dimension,
-{
-    fn op_multi(&self, a: &ArrayBase<S, D>) -> Array<A, D>;
-}
+    /// Apply operator with move
+    fn apply_into<S>(&self, mut a: ArrayBase<S, Ix1>) -> ArrayBase<S, Ix1>
+    where
+        S: DataOwned<Elem = Self::Elem> + DataMut,
+    {
+        self.apply_mut(&mut a);
+        a
+    }
 
-impl<T, A, S, D> OperatorMulti<A, S, D> for T
-where
-    A: Scalar + Lapack,
-    S: DataMut<Elem = A>,
-    D: Dimension + RemoveAxis,
-    for<'a> T: OperatorInplace<ViewRepr<&'a mut A>, D::Smaller>,
-{
-    fn op_multi(&self, a: &ArrayBase<S, D>) -> Array<A, D> {
-        let a = a.to_owned();
-        self.op_multi_into(a)
+    /// Apply operator to matrix out-place
+    fn apply2<S>(&self, a: &ArrayBase<S, Ix2>) -> Array2<S::Elem>
+    where
+        S: Data<Elem = Self::Elem>,
+    {
+        let cols: Vec<_> = a.axis_iter(Axis(1)).map(|col| self.apply(&col)).collect();
+        hstack(&cols).unwrap()
     }
-}
 
-pub trait OperatorMultiInto<S, D>
-where
-    S: DataMut,
-    D: Dimension,
-{
-    fn op_multi_into(&self, a: ArrayBase<S, D>) -> ArrayBase<S, D>;
-}
+    /// Apply operator to matrix in-place
+    fn apply2_mut<S>(&self, a: &mut ArrayBase<S, Ix2>)
+    where
+        S: DataMut<Elem = Self::Elem>,
+    {
+        for mut col in a.axis_iter_mut(Axis(1)) {
+            self.apply_mut(&mut col)
+        }
+    }
 
-impl<T, A, S, D> OperatorMultiInto<S, D> for T
-where
-    S: DataMut<Elem = A>,
-    D: Dimension + RemoveAxis,
-    for<'a> T: OperatorInplace<ViewRepr<&'a mut A>, D::Smaller>,
-{
-    fn op_multi_into(&self, mut a: ArrayBase<S, D>) -> ArrayBase<S, D> {
-        self.op_multi_inplace(&mut a);
+    /// Apply operator to matrix with move
+    fn apply2_into<S>(&self, mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
+    where
+        S: DataOwned<Elem = Self::Elem> + DataMut,
+    {
+        self.apply2_mut(&mut a);
         a
     }
 }
 
-pub trait OperatorMultiInplace<S, D>
+impl<A, Sa> LinearOperator for ArrayBase<Sa, Ix2>
 where
-    S: DataMut,
-    D: Dimension,
+    A: Scalar,
+    Sa: Data<Elem = A>,
 {
-    fn op_multi_inplace<'a>(&self, a: &'a mut ArrayBase<S, D>) -> &'a mut ArrayBase<S, D>;
-}
+    type Elem = A;
 
-impl<T, A, S, D> OperatorMultiInplace<S, D> for T
-where
-    S: DataMut<Elem = A>,
-    D: Dimension + RemoveAxis,
-    for<'a> T: OperatorInplace<ViewRepr<&'a mut A>, D::Smaller>,
-{
-    fn op_multi_inplace<'a>(&self, a: &'a mut ArrayBase<S, D>) -> &'a mut ArrayBase<S, D> {
-        let n = a.ndim();
-        for mut col in a.axis_iter_mut(Axis(n - 1)) {
-            self.op_inplace(&mut col);
-        }
-        a
+    fn apply<S>(&self, a: &ArrayBase<S, Ix1>) -> Array1<A>
+    where
+        S: Data<Elem = A>,
+    {
+        self.dot(a)
     }
 }

src/operator.rs

@@ -5,15 +5,15 @@ use ndarray_linalg::*;
 fn diag_1d() {
     let d = arr1(&[1.0, 2.0]);
     let v = arr1(&[1.0, 1.0]);
-    let dv = d.into_diagonal().op(&v);
+    let dv = d.into_diagonal().apply(&v);
     assert_close_l2!(&dv, &arr1(&[1.0, 2.0]), 1e-7);
 }
 
 #[test]
 fn diag_2d() {
     let d = arr1(&[1.0, 2.0]);
     let m = arr2(&[[1.0, 1.0], [1.0, 1.0]]);
-    let dm = d.into_diagonal().op(&m);
+    let dm = d.into_diagonal().apply2(&m);
     println!("dm = {:?}", dm);
     assert_close_l2!(&dm, &arr2(&[[1.0, 1.0], [2.0, 2.0]]), 1e-7);
 }
@@ -22,7 +22,7 @@ fn diag_2d() {
 fn diag_2d_multi() {
     let d = arr1(&[1.0, 2.0]);
     let m = arr2(&[[1.0, 1.0], [1.0, 1.0]]);
-    let dm = d.into_diagonal().op_multi_into(m);
+    let dm = d.into_diagonal().apply2_into(m);
     println!("dm = {:?}", dm);
     assert_close_l2!(&dm, &arr2(&[[1.0, 1.0], [2.0, 2.0]]), 1e-7);
 }