From a03fd6bff7e5670ad49dfcc77ab9926a3611f065 Mon Sep 17 00:00:00 2001 From: sebcrozet Date: Sat, 22 Sep 2018 15:38:33 +0200 Subject: [PATCH] Add zip_zip_map to map on three matrices simultaneously. --- src/base/matrix.rs | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/src/base/matrix.rs b/src/base/matrix.rs index 0fff3efc..ee656bdc 100644 --- a/src/base/matrix.rs +++ b/src/base/matrix.rs @@ -393,6 +393,42 @@ impl> Matrix { res } + /// Returns a matrix containing the result of `f` applied to each entries of `self` and + /// `b`, and `c`. + #[inline] + pub fn zip_zip_map(&self, b: &Matrix, c: &Matrix, mut f: F) -> MatrixMN + where + N2: Scalar, + N3: Scalar, + N4: Scalar, + S2: Storage, + S3: Storage, + F: FnMut(N, N2, N3) -> N4, + DefaultAllocator: Allocator, + { + let (nrows, ncols) = self.data.shape(); + + let mut res = unsafe { MatrixMN::new_uninitialized_generic(nrows, ncols) }; + + assert!( + (nrows.value(), ncols.value()) == b.shape() && (nrows.value(), ncols.value()) == c.shape(), + "Matrix simultaneous traversal error: dimension mismatch." + ); + + for j in 0..ncols.value() { + for i in 0..nrows.value() { + unsafe { + let a = *self.data.get_unchecked(i, j); + let b = *b.data.get_unchecked(i, j); + let c = *c.data.get_unchecked(i, j); + *res.data.get_unchecked_mut(i, j) = f(a, b, c) + } + } + } + + res + } + /// Transposes `self` and store the result into `out`. #[inline] pub fn transpose_to(&self, out: &mut Matrix)