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