ndarray_linalg/
operator.rs

1//! Linear operator algebra
2
3use crate::generate::hstack;
4use crate::types::*;
5use ndarray::*;
6
7/// Abstracted linear operator as an action to vector (`ArrayBase<S, Ix1>`) and matrix
8/// (`ArrayBase<S, Ix2`)
9pub trait LinearOperator {
10    type Elem: Scalar;
11
12    /// Apply operator out-place
13    fn apply(&self, a: &ArrayRef<Self::Elem, Ix1>) -> Array1<Self::Elem> {
14        let mut a = a.to_owned();
15        self.apply_mut(&mut a);
16        a
17    }
18
19    /// Apply operator in-place
20    fn apply_mut(&self, a: &mut ArrayRef<Self::Elem, Ix1>) {
21        let b = self.apply(a);
22        azip!((a in a, &b in &b) *a = b);
23    }
24
25    /// Apply operator with move
26    fn apply_into<S>(&self, mut a: ArrayBase<S, Ix1>) -> ArrayBase<S, Ix1>
27    where
28        S: DataOwned<Elem = Self::Elem> + DataMut,
29    {
30        self.apply_mut(&mut a);
31        a
32    }
33
34    /// Apply operator to matrix out-place
35    fn apply2(&self, a: &ArrayRef<Self::Elem, Ix2>) -> Array2<Self::Elem> {
36        let cols: Vec<_> = a.axis_iter(Axis(1)).map(|col| self.apply(&col)).collect();
37        hstack(&cols).unwrap()
38    }
39
40    /// Apply operator to matrix in-place
41    fn apply2_mut(&self, a: &mut ArrayRef<Self::Elem, Ix2>) {
42        for mut col in a.axis_iter_mut(Axis(1)) {
43            self.apply_mut(&mut col)
44        }
45    }
46
47    /// Apply operator to matrix with move
48    fn apply2_into<S>(&self, mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
49    where
50        S: DataOwned<Elem = Self::Elem> + DataMut,
51    {
52        self.apply2_mut(&mut a);
53        a
54    }
55}
56
57impl<A, Sa> LinearOperator for ArrayBase<Sa, Ix2>
58where
59    A: Scalar,
60    Sa: Data<Elem = A>,
61{
62    type Elem = A;
63
64    fn apply(&self, a: &ArrayRef<A, Ix1>) -> Array1<A> {
65        self.dot(a)
66    }
67}