Skip to content

Commit 316e284

Browse files
committed
Add impl for ArrayBase
1 parent ece39bb commit 316e284

File tree

1 file changed

+28
-12
lines changed

1 file changed

+28
-12
lines changed

src/krylov/operator.rs

Lines changed: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,12 @@ use crate::types::*;
55
use ndarray::*;
66

77
pub trait LinearOperator {
8+
type Elem: Scalar;
9+
810
/// Apply operator out-place
911
fn apply<S>(&self, a: &ArrayBase<S, Ix1>) -> Array1<S::Elem>
1012
where
11-
S: Data,
12-
S::Elem: Scalar,
13+
S: Data<Elem = Self::Elem>,
1314
{
1415
let mut a = a.to_owned();
1516
self.apply_mut(&mut a);
@@ -18,13 +19,16 @@ pub trait LinearOperator {
1819
/// Apply operator in-place
1920
fn apply_mut<S>(&self, a: &mut ArrayBase<S, Ix1>)
2021
where
21-
S: DataMut,
22-
S::Elem: Scalar;
22+
S: DataMut<Elem = Self::Elem>,
23+
{
24+
let b = self.apply(a);
25+
azip!(mut a(a), b in { *a = b });
26+
}
27+
2328
/// Apply operator with move
2429
fn apply_into<S>(&self, mut a: ArrayBase<S, Ix1>) -> ArrayBase<S, Ix1>
2530
where
26-
S: DataOwned + DataMut,
27-
S::Elem: Scalar,
31+
S: DataOwned<Elem = Self::Elem> + DataMut,
2832
{
2933
self.apply_mut(&mut a);
3034
a
@@ -33,17 +37,15 @@ pub trait LinearOperator {
3337
/// Apply operator to matrix out-place
3438
fn apply2<S>(&self, a: &ArrayBase<S, Ix2>) -> Array2<S::Elem>
3539
where
36-
S: Data,
37-
S::Elem: Scalar,
40+
S: Data<Elem = Self::Elem>,
3841
{
3942
let cols: Vec<_> = a.axis_iter(Axis(0)).map(|col| self.apply(&col)).collect();
4043
hstack(&cols).unwrap()
4144
}
4245
/// Apply operator to matrix in-place
4346
fn apply2_mut<S>(&self, a: &mut ArrayBase<S, Ix2>)
4447
where
45-
S: DataMut,
46-
S::Elem: Scalar,
48+
S: DataMut<Elem = Self::Elem>,
4749
{
4850
for mut col in a.axis_iter_mut(Axis(0)) {
4951
self.apply_mut(&mut col)
@@ -52,10 +54,24 @@ pub trait LinearOperator {
5254
/// Apply operator to matrix with move
5355
fn apply2_into<S>(&self, mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>
5456
where
55-
S: DataOwned + DataMut,
56-
S::Elem: Scalar,
57+
S: DataOwned<Elem = Self::Elem> + DataMut,
5758
{
5859
self.apply2_mut(&mut a);
5960
a
6061
}
6162
}
63+
64+
impl<A, Sa> LinearOperator for ArrayBase<Sa, Ix2>
65+
where
66+
A: Scalar,
67+
Sa: Data<Elem = A>,
68+
{
69+
type Elem = A;
70+
71+
fn apply<S>(&self, a: &ArrayBase<S, Ix1>) -> Array1<A>
72+
where
73+
S: Data<Elem = A>,
74+
{
75+
self.dot(a)
76+
}
77+
}

0 commit comments

Comments
 (0)