Generalize colzip and colzipwith to any matrices.

Add zip2three to zip three matrices together.

diff --git a/src/Linear/Matrix.hs b/src/Linear/Matrix.hs
index c33e96691b00f9cb2c87c3332ab572a32286d10b..3d4daab3dc1f7be81277238b52253c5df54ae61f 100644 (file)
--- a/src/Linear/Matrix.hs
+++ b/src/Linear/Matrix.hs
@@ -818,37 +818,73 @@ trace matrix =
     element_sum $ V.map V.head rows
 
 
--- | Zip together two column matrices.
+-- | Zip together two matrices.
+--
+--   TODO: don't cheat with construct (map V.zips instead).
 --
 --   Examples:
 --
 --   >>> let m1 = fromList [[1],[1],[1]] :: Col3 Int
 --   >>> let m2 = fromList [[1],[2],[3]] :: Col3 Int
---   >>> colzip m1 m2
+--   >>> zip2 m1 m2
 --   (((1,1)),((1,2)),((1,3)))
 --
-colzip :: Arity m => Col m a -> Col m a -> Col m (a,a)
-colzip c1 c2 =
+--   >>> let m1 = fromList [[1,2],[3,4]] :: Mat2 Int
+--   >>> let m2 = fromList [[1,1],[1,1]] :: Mat2 Int
+--   >>> zip2 m1 m2
+--   (((1,1),(2,1)),((3,1),(4,1)))
+--
+zip2 :: (Arity m, Arity n) => Mat m n a -> Mat m n a -> Mat m n (a,a)
+zip2 m1 m2 =
+  construct lambda
+  where
+    lambda i j = (m1 !!! (i,j), m2 !!! (i,j))
+
+
+-- | Zip together three matrices.
+--
+--   TODO: don't cheat with construct (map V.zips instead).
+--
+--   Examples:
+--
+--   >>> let m1 = fromList [[1],[1],[1]] :: Col3 Int
+--   >>> let m2 = fromList [[1],[2],[3]] :: Col3 Int
+--   >>> let m3 = fromList [[4],[5],[6]] :: Col3 Int
+--   >>> zip2three m1 m2 m3
+--   (((1,1,4)),((1,2,5)),((1,3,6)))
+--
+--   >>> let m1 = fromList [[1,2],[3,4]] :: Mat2 Int
+--   >>> let m2 = fromList [[1,1],[1,1]] :: Mat2 Int
+--   >>> let m3 = fromList [[8,2],[6,3]] :: Mat2 Int
+--   >>> zip2three m1 m2 m3
+--   (((1,1,8),(2,1,2)),((3,1,6),(4,1,3)))
+--
+zip2three :: (Arity m, Arity n)
+          => Mat m n a
+          -> Mat m n a
+          -> Mat m n a
+          -> Mat m n (a,a,a)
+zip2three m1 m2 m3 =
   construct lambda
   where
-    lambda i j = (c1 !!! (i,j), c2 !!! (i,j))
+    lambda i j = (m1 !!! (i,j), m2 !!! (i,j), m3 !!! (i,j))
 
 
--- | Zip together two column matrices using the supplied function.
+-- | Zip together two matrices using the supplied function.
 --
 --   Examples:
 --
 --   >>> let c1 = fromList [[1],[2],[3]] :: Col3 Integer
 --   >>> let c2 = fromList [[4],[5],[6]] :: Col3 Integer
---   >>> colzipwith (^) c1 c2
+--   >>> zipwith2 (^) c1 c2
 --   ((1),(32),(729))
 --
-colzipwith :: Arity m
+zipwith2 :: Arity m
            => (a -> a -> b)
            -> Col m a
            -> Col m a
            -> Col m b
-colzipwith f c1 c2 =
+zipwith2 f c1 c2 =
   construct lambda
   where
     lambda i j = f (c1 !!! (i,j)) (c2 !!! (i,j))
@@ -931,3 +967,5 @@ imap2 f (Mat rows) =
 --
 reverse2 :: (Arity m, Arity n) => Mat m n a -> Mat m n a
 reverse2 (Mat rows) = Mat $ V.reverse $ V.map V.reverse rows
+
+