Generalize colzip and colzipwith to any matrices.

[numerical-analysis.git] / src / Linear / Matrix.hs

diff --git a/src/Linear/Matrix.hs b/src/Linear/Matrix.hs
index e4acc9ae34818f732310129279d13cb503f96c0e..3d4daab3dc1f7be81277238b52253c5df54ae61f 100644 (file)
--- a/src/Linear/Matrix.hs
+++ b/src/Linear/Matrix.hs
@@ -31,8 +31,7 @@ import Data.Vector.Fixed (
   mk2,
   mk3,
   mk4,
-  mk5
-  )
+  mk5 )
 import qualified Data.Vector.Fixed as V (
   and,
   fromList,
@@ -42,6 +41,7 @@ import qualified Data.Vector.Fixed as V (
   map,
   maximum,
   replicate,
+  reverse,
   toList,
   zipWith )
 import Data.Vector.Fixed.Cont ( Arity, arity )
@@ -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))
@@ -875,7 +911,7 @@ map2 f (Mat rows) =
 --   Examples:
 --
 --   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
---   >>>  ifoldl2 (\i j cur _ -> cur + i + j) 0 m
+--   >>> ifoldl2 (\i j cur _ -> cur + i + j) 0 m
 --   18
 --
 ifoldl2 :: forall a b m n.
@@ -911,3 +947,25 @@ imap2 f (Mat rows) =
   Mat $ V.imap g rows
   where
     g i = V.imap (f i)
+
+
+-- | Reverse the order of elements in a matrix.
+--
+--   Examples:
+--
+--   >>> let m1 = fromList [[1,2,3]] :: Row3 Int
+--   >>> reverse2 m1
+--   ((3,2,1))
+--
+--   >>> let m1 = vec3d (1,2,3 :: Int)
+--   >>> reverse2 m1
+--   ((3),(2),(1))
+--
+--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
+--   >>> reverse2 m
+--   ((9,8,7),(6,5,4),(3,2,1))
+--
+reverse2 :: (Arity m, Arity n) => Mat m n a -> Mat m n a
+reverse2 (Mat rows) = Mat $ V.reverse $ V.map V.reverse rows
+
+