Skip to content

Commit

Permalink
Leverage DenseVector in permutation * dense product
Browse files Browse the repository at this point in the history
Now it's possible to call on ndarray data as well, with a single impl.
  • Loading branch information
vbarrielle committed Jan 27, 2021
1 parent 79bd165 commit 210bd1d
Showing 1 changed file with 16 additions and 8 deletions.
24 changes: 16 additions & 8 deletions src/sparse/permutation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
/// Both the permutation matrices and its inverse are stored
use std::ops::{Deref, Mul};

use crate::dense_vector::{DenseVector, DenseVectorMut};
use crate::indexing::SpIndex;
use crate::sparse::{CompressedStorage, CsMatI, CsMatViewI};

Expand Down Expand Up @@ -251,21 +252,24 @@ where
}
}

impl<'a, 'b, N, I, IndStorage> Mul<&'a [N]> for &'b Permutation<I, IndStorage>
impl<'b, V, I, IndStorage> Mul<V> for &'b Permutation<I, IndStorage>
where
IndStorage: 'b + Deref<Target = [I]>,
N: 'a + Copy,
V: DenseVector,
<V as DenseVector>::Owned:
DenseVectorMut + DenseVector<Scalar = <V as DenseVector>::Scalar>,
<V as DenseVector>::Scalar: Clone,
I: SpIndex,
{
type Output = Vec<N>;
fn mul(self, rhs: &'a [N]) -> Vec<N> {
assert_eq!(self.dim, rhs.len());
let mut res = rhs.to_vec();
type Output = V::Owned;
fn mul(self, rhs: V) -> Self::Output {
assert_eq!(self.dim, rhs.dim());
let mut res = rhs.to_owned();
match self.storage {
Identity => res,
FinitePerm { perm: ref p, .. } => {
for (pi, r) in p.iter().zip(res.iter_mut()) {
*r = rhs[pi.index_unchecked()];
for (i, pi) in p.iter().enumerate() {
*res.index_mut(i) = rhs.index(pi.index_unchecked()).clone();
}
res
}
Expand Down Expand Up @@ -344,6 +348,10 @@ mod test {

let y = &p * &x;
assert_eq!(&y, &[2, 1, 3, 5, 4]);

let x = ndarray::arr1(&[5, 1, 2, 3, 4]);
let y = &p * x.view();
assert_eq!(y, ndarray::arr1(&[2, 1, 3, 5, 4]));
}

#[test]
Expand Down

0 comments on commit 210bd1d

Please sign in to comment.