Implement map_with_location.

src/base/matrix.rs

@@ -360,6 +360,29 @@ impl<N: Scalar, R: Dim, C: Dim, S: Storage<N, R, C>> Matrix<N, R, C, S> {
         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: Scalar, F: FnMut(usize, usize, N) -> N2>(&self, mut f: F) -> MatrixMN<N2, R, C>
+    where
+        DefaultAllocator: Allocator<N2, R, C>,
+    {
+        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]

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);