diff --git a/src/base/matrix.rs b/src/base/matrix.rs index 625438d5..3d12c12b 100644 --- a/src/base/matrix.rs +++ b/src/base/matrix.rs @@ -706,8 +706,7 @@ impl> Matrix { /// Replaces each component of `self` by the result of a closure `f` applied on it. #[inline] - pub fn apply N>(&mut self, mut f: F) - where DefaultAllocator: Allocator { + pub fn apply N>(&mut self, mut f: F) { let (nrows, ncols) = self.shape(); for j in 0..ncols { @@ -719,6 +718,71 @@ impl> Matrix { } } } + + /// Replaces each component of `self` by the result of a closure `f` applied on its components + /// joined with the components from `rhs`. + #[inline] + pub fn zip_apply(&mut self, rhs: Matrix, mut f: impl FnMut(N, N2) -> N) + where N2: Scalar, + R2: Dim, + C2: Dim, + S2: Storage, + ShapeConstraint: SameNumberOfRows + SameNumberOfColumns { + let (nrows, ncols) = self.shape(); + + assert!( + (nrows, ncols) == rhs.shape(), + "Matrix simultaneous traversal error: dimension mismatch." + ); + + for j in 0..ncols { + for i in 0..nrows { + unsafe { + let e = self.data.get_unchecked_mut(i, j); + let rhs = rhs.get_unchecked(i, j); + *e = f(*e, *rhs) + } + } + } + } + + + /// Replaces each component of `self` by the result of a closure `f` applied on its components + /// joined with the components from `b` and `c`. + #[inline] + pub fn zip_zip_apply(&mut self, b: Matrix, c: Matrix, mut f: impl FnMut(N, N2, N3) -> N) + where N2: Scalar, + R2: Dim, + C2: Dim, + S2: Storage, + N3: Scalar, + R3: Dim, + C3: Dim, + S3: Storage, + ShapeConstraint: SameNumberOfRows + SameNumberOfColumns, + ShapeConstraint: SameNumberOfRows + SameNumberOfColumns { + let (nrows, ncols) = self.shape(); + + assert!( + (nrows, ncols) == b.shape(), + "Matrix simultaneous traversal error: dimension mismatch." + ); + assert!( + (nrows, ncols) == c.shape(), + "Matrix simultaneous traversal error: dimension mismatch." + ); + + for j in 0..ncols { + for i in 0..nrows { + unsafe { + let e = self.data.get_unchecked_mut(i, j); + let b = b.get_unchecked(i, j); + let c = c.get_unchecked(i, j); + *e = f(*e, *b, *c) + } + } + } + } } impl> Vector {