ndarray_linalg/
operator.rs1use crate::generate::hstack;
4use crate::types::*;
5use ndarray::*;
6
7pub trait LinearOperator {
10 type Elem: Scalar;
11
12 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 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 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 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 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 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}