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<S>(&self, a: &ArrayBase<S, Ix1>) -> Array1<S::Elem>
14    where
15        S: Data<Elem = Self::Elem>,
16    {
17        let mut a = a.to_owned();
18        self.apply_mut(&mut a);
19        a
20    }
21
22    /// Apply operator in-place
23    fn apply_mut<S>(&self, a: &mut ArrayBase<S, Ix1>)
24    where
25        S: DataMut<Elem = Self::Elem>,
26    {
27        let b = self.apply(a);
28        azip!((a in a, &b in &b) *a = b);
29    }
30
31    /// Apply operator with move
32    fn apply_into<S>(&self, mut a: ArrayBase<S, Ix1>) -> ArrayBase<S, Ix1>
33    where
34        S: DataOwned<Elem = Self::Elem> + DataMut,
35    {
36        self.apply_mut(&mut a);
37        a
38    }
39
40    /// Apply operator to matrix out-place
41    fn apply2<S>(&self, a: &ArrayBase<S, Ix2>) -> Array2<S::Elem>
42    where
43        S: Data<Elem = Self::Elem>,
44    {
45        let cols: Vec<_> = a.axis_iter(Axis(1)).map(|col| self.apply(&col)).collect();
46        hstack(&cols).unwrap()
47    }
48
49    /// Apply operator to matrix in-place
50    fn apply2_mut<S>(&self, a: &mut ArrayBase<S, Ix2>)
51    where
52        S: DataMut<Elem = Self::Elem>,
53    {
54        for mut col in a.axis_iter_mut(Axis(1)) {
55            self.apply_mut(&mut col)
56        }
57    }
58
59    /// Apply operator to matrix with move
60    fn apply2_into<S>(&self, mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
61    where
62        S: DataOwned<Elem = Self::Elem> + DataMut,
63    {
64        self.apply2_mut(&mut a);
65        a
66    }
67}
68
69impl<A, Sa> LinearOperator for ArrayBase<Sa, Ix2>
70where
71    A: Scalar,
72    Sa: Data<Elem = A>,
73{
74    type Elem = A;
75
76    fn apply<S>(&self, a: &ArrayBase<S, Ix1>) -> Array1<A>
77    where
78        S: Data<Elem = A>,
79    {
80        self.dot(a)
81    }
82}