Add the imap2 function to Linear.Matrix.

diff --git a/src/Linear/Matrix.hs b/src/Linear/Matrix.hs
index 611545fc9dc44e97bdf3c30f0333c7a7b7064ce6..cfa838042ee7bf0794e7a5275981f7ab16ffd7fb 100644 (file)
--- a/src/Linear/Matrix.hs
+++ b/src/Linear/Matrix.hs
@@ -38,6 +38,7 @@ import qualified Data.Vector.Fixed as V (
   fromList,
   head,
   ifoldl,
+  imap,
   map,
   maximum,
   replicate,
@@ -871,3 +872,19 @@ ifoldl2 f initial (Mat rows) =
 
     row_function :: b -> Int -> Vec n a -> b
     row_function rowinit idx r = V.ifoldl (g idx) rowinit r
+
+
+-- | Map a function over a matrix of any dimensions, passing the
+--   coordinates @i@ and @j@ to the function @f@.
+--
+--   Examples:
+--
+--   >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
+--   >>> imap2 (\i j _ -> i+j) m
+--   ((0,1),(1,2))
+--
+imap2 :: (Int -> Int -> a -> b) -> Mat m n a -> Mat m n b
+imap2 f (Mat rows) =
+  Mat $ V.imap g rows
+  where
+    g i = V.imap (f i)