toBlocks

3 files changed, 41 insertions, 3 deletions

diff --git a/lib/Data/Packed/Internal/Vector.hs b/lib/Data/Packed/Internal/Vector.hs
index 04232a2..9ca8e58 100644
--- a/lib/Data/Packed/Internal/Vector.hs
+++ b/lib/Data/Packed/Internal/Vector.hs
@@ -16,7 +16,7 @@
 module Data.Packed.Internal.Vector (
     Vector, dim,
     fromList, toList, (|>),
-    join, (@>), safe, at, at', subVector,
+    join, (@>), safe, at, at', subVector, takesV,
     mapVector, zipVector,
     foldVector, foldVectorG, foldLoop,
     createVector, vec,
@@ -222,6 +222,21 @@ join as = unsafePerformIO $ do
             joiner cs 0 (advancePtr p n)
 
 
+{- | Extract consecutive subvectors of the given sizes.
+
+@> takesV [3,4] (linspace 10 (1,10))
+[3 |> [1.0,2.0,3.0],4 |> [4.0,5.0,6.0,7.0]]@
+
+-}
+takesV :: Storable t => [Int] -> Vector t -> [Vector t]
+takesV ms w | sum ms > dim w = error $ "takesV " ++ show ms ++ " on dim = " ++ (show $ dim w)
+            | otherwise = go ms w
+    where go [] _ = []
+          go (n:ns) v = subVector 0 n v
+                      : go ns (subVector n (dim v - n) v)
+
+---------------------------------------------------------------
+
 -- | transforms a complex vector into a real vector with alternating real and imaginary parts 
 asReal :: Vector (Complex Double) -> Vector Double
 asReal v = V { ioff = 2*ioff v, idim = 2*dim v, fptr =  castForeignPtr (fptr v) }
diff --git a/lib/Data/Packed/Matrix.hs b/lib/Data/Packed/Matrix.hs
index e204a83..d91a089 100644
--- a/lib/Data/Packed/Matrix.hs
+++ b/lib/Data/Packed/Matrix.hs
@@ -22,7 +22,8 @@ module Data.Packed.Matrix (
     (@@>),
     asRow, asColumn,
     fromRows, toRows, fromColumns, toColumns,
-    fromBlocks, repmat,
+    fromBlocks, toBlocks, toBlocksEvery,
+    repmat,
     flipud, fliprl,
     subMatrix, takeRows, dropRows, takeColumns, dropColumns,
     extractRows,
@@ -413,3 +414,25 @@ compat' m1 m2 = rows m1 == 1 && cols m1 == 1
              || rows m2 == 1 && cols m2 == 1
              || rows m1 == rows m2 && cols m1 == cols m2
 
+------------------------------------------------------------
+
+toBlockRows [r] m | r == rows m = [m]
+toBlockRows rs m = map (reshape (cols m)) (takesV szs (flatten m))
+    where szs = map (* cols m) rs
+
+toBlockCols [c] m | c == cols m = [m]
+toBlockCols cs m = map trans . toBlockRows cs . trans $ m
+
+-- | Partition a matrix into blocks with the given numbers of rows and columns.
+-- The remaining rows and columns are discarded.
+toBlocks :: (Element t) => [Int] -> [Int] -> Matrix t -> [[Matrix t]]
+toBlocks rs cs m = map (toBlockCols cs) . toBlockRows rs $ m
+
+-- | Fully partition a matrix into blocks of the same size. If the dimensions are not
+-- a multiple of the given size the last blocks will be smaller.
+toBlocksEvery :: (Element t) => Int -> Int -> Matrix t -> [[Matrix t]]
+toBlocksEvery r c m = toBlocks rs cs m where
+    (qr,rr) = rows m `divMod` r
+    (qc,rc) = cols m `divMod` c
+    rs = replicate qr r ++ if rr > 0 then [rr] else []
+    cs = replicate qc c ++ if rc > 0 then [rc] else []
diff --git a/lib/Data/Packed/Vector.hs b/lib/Data/Packed/Vector.hs
index 6f1096f..472240c 100644
--- a/lib/Data/Packed/Vector.hs
+++ b/lib/Data/Packed/Vector.hs
@@ -16,7 +16,7 @@ module Data.Packed.Vector (
     Vector,
     fromList, (|>), toList, buildVector,
     dim, (@>),
-    subVector, join,
+    subVector, takesV, join,
     constant, linspace,
     vecdisp,
     vectorMax, vectorMin, vectorMaxIndex, vectorMinIndex,