From 12962c3c134e32b01f192aeebd6264a045ec43b0 Mon Sep 17 00:00:00 2001 From: mborst Date: Tue, 25 Sep 2018 21:27:27 +0200 Subject: [PATCH] Implement map_with_location. --- src/base/matrix.rs | 23 +++++++++++++++++++++++ tests/core/matrix.rs | 11 +++++++++++ 2 files changed, 34 insertions(+) diff --git a/src/base/matrix.rs b/src/base/matrix.rs index 0336df71..ebd35b0b 100644 --- a/src/base/matrix.rs +++ b/src/base/matrix.rs @@ -360,6 +360,29 @@ impl> Matrix { res } + /// Returns a matrix containing the result of `f` applied to each of its entries. Unlike `map`, + /// `f` also gets passed the row and column index, i.e. `f(value, row, col)`. + #[inline] + pub fn map_with_location N2>(&self, mut f: F) -> MatrixMN + where + DefaultAllocator: Allocator, + { + let (nrows, ncols) = self.data.shape(); + + let mut res = unsafe { MatrixMN::new_uninitialized_generic(nrows, ncols) }; + + for j in 0..ncols.value() { + for i in 0..nrows.value() { + unsafe { + let a = *self.data.get_unchecked(i, j); + *res.data.get_unchecked_mut(i, j) = f(i, j, a) + } + } + } + + res + } + /// Returns a matrix containing the result of `f` applied to each entries of `self` and /// `rhs`. #[inline] diff --git a/tests/core/matrix.rs b/tests/core/matrix.rs index 34ecaad7..50a2f860 100644 --- a/tests/core/matrix.rs +++ b/tests/core/matrix.rs @@ -425,6 +425,17 @@ fn map() { assert_eq!(computed, expected); } +#[test] +fn map_with_location() { + let a = Matrix4::new(1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4); + + let expected = Matrix4::new(1, 2, 3, 4, 3, 4, 5, 6, 5, 6, 7, 8, 7, 8, 9, 10); + + let computed = a.map_with_location(|i, j, e| e + i + j); + + assert_eq!(computed, expected); +} + #[test] fn zip_map() { let a = Matrix3::new(11i32, 12, 13, 21, 22, 23, 31, 32, 33);