1 files changed, 29 insertions, 15 deletions

diff --git a/lib/Data/Packed/Matrix.hs b/lib/Data/Packed/Matrix.hs
index c1a9b24..1b67820 100644
--- a/lib/Data/Packed/Matrix.hs
+++ b/lib/Data/Packed/Matrix.hs
@@ -30,7 +30,7 @@ module Data.Packed.Matrix (
     (@@>),
     asRow, asColumn,
     fromRows, toRows, fromColumns, toColumns,
-    fromBlocks, toBlocks, toBlocksEvery,
+    fromBlocks, diagBlock, toBlocks, toBlocksEvery,
     repmat,
     flipud, fliprl,
     subMatrix, takeRows, dropRows, takeColumns, dropColumns,
@@ -46,6 +46,7 @@ import Data.Array
 
 import Data.List(transpose,intersperse)
 import Foreign.Storable(Storable)
+import Control.Monad(liftM)
 
 -------------------------------------------------------------------
 
@@ -155,7 +156,19 @@ adaptBlocks ms = ms' where
         x = m@@>(0,0)
     g _ _ = error "inconsistent dimensions in fromBlocks"
 
------------------------------------------------------------
+
+--------------------------------------------------------------------------------
+
+-- | create a block diagonal matrix
+diagBlock :: (Element t, Num t) => [Matrix t] -> Matrix t
+diagBlock ms = fromBlocks $ zipWith f ms [0..]
+  where
+    f m k = take n $ replicate k z ++ m : repeat z
+    n = length ms
+    z = (1><1) [0]
+
+--------------------------------------------------------------------------------
+
 
 -- | Reverse rows
 flipud :: Element t => Matrix t -> Matrix t
@@ -351,7 +364,11 @@ toBlocksEvery r c m = toBlocks rs cs m where
 
 -------------------------------------------------------------------
 
-mk c g = \k v -> g (divMod k c) v
+-- Given a column number and a function taking matrix indexes, returns
+-- a function which takes vector indexes (that can be used on the
+-- flattened matrix).
+mk :: Int -> ((Int, Int) -> t) -> (Int -> t)
+mk c g = \k -> g (divMod k c)
 
 {- |
 
@@ -364,9 +381,8 @@ m[1,1] = 5
 m[1,2] = 6@
 -}
 mapMatrixWithIndexM_
-  :: (Element a, Num a,
-      Functor f, Monad f) =>
-      ((Int, Int) -> a -> f ()) -> Matrix a -> f ()
+  :: (Element a, Num a, Monad m) =>
+      ((Int, Int) -> a -> m ()) -> Matrix a -> m ()
 mapMatrixWithIndexM_ g m = mapVectorWithIndexM_ (mk c g) . flatten $ m
   where
     c = cols m
@@ -380,11 +396,9 @@ Just (3><3)
  ,  20.0,  21.0, 122.0 ]@
 -}
 mapMatrixWithIndexM
-  :: (Foreign.Storable.Storable t,
-      Element a, Num a,
-      Functor f, Monad f) =>
-      ((Int, Int) -> a -> f t) -> Matrix a -> f (Matrix t)
-mapMatrixWithIndexM g m = fmap (reshape c) . mapVectorWithIndexM (mk c g) . flatten $ m
+  :: (Element a, Storable b, Monad m) =>
+      ((Int, Int) -> a -> m b) -> Matrix a -> m (Matrix b)
+mapMatrixWithIndexM g m = liftM (reshape c) . mapVectorWithIndexM (mk c g) . flatten $ m 
     where
       c = cols m
 
@@ -395,10 +409,10 @@ mapMatrixWithIndexM g m = fmap (reshape c) . mapVectorWithIndexM (mk c g) . flat
  ,  10.0, 111.0,  12.0
  ,  20.0,  21.0, 122.0 ]@
  -}
-mapMatrixWithIndex :: (Foreign.Storable.Storable t,
-      Element a, Num a) =>
-      ((Int, Int) -> a -> t) -> Matrix a -> Matrix t
-mapMatrixWithIndex g m = reshape c $ mapVectorWithIndex (mk c g) $ flatten m
+mapMatrixWithIndex
+  :: (Element a, Storable b) =>
+      ((Int, Int) -> a -> b) -> Matrix a -> Matrix b
+mapMatrixWithIndex g m = reshape c . mapVectorWithIndex (mk c g) . flatten $ m
     where
       c = cols m