more defined fromRows/fromColumns/fromBlocks

3 files changed, 60 insertions, 18 deletions

diff --git a/lib/Data/Packed/Internal/Common.hs b/lib/Data/Packed/Internal/Common.hs
index 629016d..972cd7d 100644
--- a/lib/Data/Packed/Internal/Common.hs
+++ b/lib/Data/Packed/Internal/Common.hs
@@ -18,7 +18,7 @@ module Data.Packed.Internal.Common(
   Adapt,
   app1, app2, app3, app4,
   (//), check,
-  partit, common,
+  partit, common, compatdim,
   fi
 ) where
 
@@ -41,6 +41,12 @@ common f = commonval . map f where
     commonval [a] = Just a
     commonval (a:b:xs) = if a==b then commonval (b:xs) else Nothing
 
+-- | common value with \"adaptable\" 1
+compatdim :: [Int] -> Maybe Int
+compatdim [] = Nothing
+compatdim [a] = Just a
+compatdim (a:b:xs) = if a==b || a==1 || b==1 then compatdim (max a b:xs) else Nothing
+
 -- | postfix function application (@flip ($)@)
 (//) :: x -> (x -> y) -> y
 infixl 0 //
diff --git a/lib/Data/Packed/Internal/Matrix.hs b/lib/Data/Packed/Internal/Matrix.hs
index a220f1a..33c9324 100644
--- a/lib/Data/Packed/Internal/Matrix.hs
+++ b/lib/Data/Packed/Internal/Matrix.hs
@@ -135,11 +135,16 @@ type Mt t s = Int -> Int -> Ptr t -> s
 toLists :: (Element t) => Matrix t -> [[t]]
 toLists m = partit (cols m) . toList . flatten $ m
 
--- | creates a Matrix from a list of vectors
+-- | Create a matrix from a list of vectors.
+-- All vectors must have the same dimension,
+-- or dimension 1, which is are automatically expanded.
 fromRows :: Element t => [Vector t] -> Matrix t
-fromRows vs = case common dim vs of
+fromRows vs = case compatdim (map dim vs) of
     Nothing -> error "fromRows applied to [] or to vectors with different sizes"
-    Just c  -> reshape c (join vs)
+    Just c  -> reshape c . join . map (adapt c) $ vs
+  where
+    adapt c v | dim v == c = v
+              | otherwise = constantD (v@>0) c
 
 -- | extracts the rows of a matrix as a list of vectors
 toRows :: Element t => Matrix t -> [Vector t]
diff --git a/lib/Data/Packed/Matrix.hs b/lib/Data/Packed/Matrix.hs
index bd7cb69..912e1c6 100644
--- a/lib/Data/Packed/Matrix.hs
+++ b/lib/Data/Packed/Matrix.hs
@@ -44,28 +44,59 @@ import Text.Printf(printf)
 -- | creates a matrix from a vertical list of matrices
 joinVert :: Element t => [Matrix t] -> Matrix t
 joinVert ms = case common cols ms of
-    Nothing -> error "joinVert on matrices with different number of columns"
+    Nothing -> error "(impossible) joinVert on matrices with different number of columns"
     Just c  -> reshape c $ join (map flatten ms)
 
 -- | creates a matrix from a horizontal list of matrices
 joinHoriz :: Element t => [Matrix t] -> Matrix t
 joinHoriz ms = trans. joinVert . map trans $ ms
 
-{- | Creates a matrix from blocks given as a list of lists of matrices:
-
-@\> let a = 'diag' $ 'fromList' [5,7,2]
-\> let b = 'reshape' 4 $ 'constant' (-1) 12
-\> fromBlocks [[a,b],[b,a]]
-(6><7)
- [  5.0,  0.0,  0.0, -1.0, -1.0, -1.0, -1.0
- ,  0.0,  7.0,  0.0, -1.0, -1.0, -1.0, -1.0
- ,  0.0,  0.0,  2.0, -1.0, -1.0, -1.0, -1.0
- , -1.0, -1.0, -1.0, -1.0,  5.0,  0.0,  0.0
- , -1.0, -1.0, -1.0, -1.0,  0.0,  7.0,  0.0
- , -1.0, -1.0, -1.0, -1.0,  0.0,  0.0,  2.0 ]@
+{- | Creates a matrix from blocks given as a list of lists of matrices.
+
+Single row/column components are automatically expanded to match the
+corresponding common row and column:
+
+@\> let disp = putStr . dispf 2
+\> let vector xs = fromList xs :: Vector Double
+\> let diagl = diag . vector
+\> let rowm = asRow . vector
+\> disp $ fromBlocks [[ident 5, 7, rowm[10,20]], [3, diagl[1,2,3], 0]]
+8x10
+1  0  0  0  0  7  7  7  10  20
+0  1  0  0  0  7  7  7  10  20
+0  0  1  0  0  7  7  7  10  20
+0  0  0  1  0  7  7  7  10  20
+0  0  0  0  1  7  7  7  10  20
+3  3  3  3  3  1  0  0   0   0
+3  3  3  3  3  0  2  0   0   0
+3  3  3  3  3  0  0  3   0   0@
 -}
 fromBlocks :: Element t => [[Matrix t]] -> Matrix t
-fromBlocks = joinVert . map joinHoriz 
+fromBlocks = fromBlocksRaw . adaptBlocks
+
+fromBlocksRaw mms = joinVert . map joinHoriz $ mms
+
+adaptBlocks ms = ms' where
+    bc = case common length ms of
+          Just c -> c
+          Nothing -> error "fromBlocks requires rectangular [[Matrix]]"
+    rs = map (compatdim . map rows) ms
+    cs = map (compatdim . map cols) (transpose ms)
+    szs = sequence [rs,cs]
+    ms' = partit bc $ zipWith g szs (concat ms)
+
+    g [Just nr,Just nc] m
+                | nr == r && nc == c = m
+                | r == 1 && c == 1 = reshape nc (constant x (nr*nc))
+                | r == 1 = fromRows (replicate nr (flatten m))
+                | otherwise = fromColumns (replicate nc (flatten m))
+      where
+        r = rows m
+        c = cols m
+        x = m@@>(0,0)
+    g _ _ = error "inconsistent dimensions in fromBlocks"
+
+-----------------------------------------------------------
 
 -- | Reverse rows 
 flipud :: Element t => Matrix t -> Matrix t