From: Michael Orlitzky <michael@orlitzky.com>
Date: Sat, 8 Feb 2014 00:33:45 +0000 (-0500)
Subject: Add the imap2 function to Linear.Matrix.
X-Git-Url: http://gitweb.michael.orlitzky.com/?p=numerical-analysis.git;a=commitdiff_plain;h=9fa2506d30661e984ff74313c3884007067387cb

Add the imap2 function to Linear.Matrix.
---

diff --git a/src/Linear/Matrix.hs b/src/Linear/Matrix.hs
index 611545f..cfa8380 100644
--- 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)