module Data.Matrix.Dense.Generic
(
Matrix(..)
, MG.dim
, MG.rows
, MG.cols
, MG.unsafeIndex
, (MG.!)
, MG.takeRow
, MG.takeColumn
, MG.takeDiag
, MG.unsafeFromVector
, MG.fromVector
, MG.matrix
, MG.fromLists
, MG.fromRows
, fromColumns
, MG.empty
, MG.flatten
, MG.toRows
, MG.toColumns
, MG.toList
, MG.toLists
, convert
, tr
, subMatrix
, ident
, diag
, diagRect
, fromBlocks
, isSymmetric
, force
, Data.Matrix.Dense.Generic.foldl
, imap
, Data.Matrix.Dense.Generic.map
, mapM
, mapM_
, forM
, forM_
, Data.Matrix.Dense.Generic.sequence
, Data.Matrix.Dense.Generic.sequence_
, generate
, MG.thaw
, MG.unsafeThaw
, MG.freeze
, MG.unsafeFreeze
, MG.create
) where
import Prelude hiding (mapM_, mapM)
import Control.Arrow ((***), (&&&))
import Control.Monad (liftM, foldM, foldM_)
import qualified Data.Foldable as F
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Matrix.Generic as MG
import Data.Matrix.Dense.Generic.Mutable (MMatrix(..))
type instance MG.Mutable Matrix = MMatrix
data Matrix v a = Matrix !Int
!Int
!Int
!Int
!(v a)
deriving (Show)
instance G.Vector v a => MG.Matrix Matrix v a where
dim (Matrix r c _ _ _) = (r,c)
unsafeIndex (Matrix _ _ tda offset vec) (i,j) = vec `G.unsafeIndex` idx
where
idx = offset + i * tda + j
unsafeFromVector (r,c) = Matrix r c c 0
unsafeTakeRow (Matrix _ c tda offset vec) i = G.slice i' c vec
where
i' = offset + i * tda
flatten (Matrix r c tda offset vec)
| c == tda = G.slice offset (r*c) vec
| otherwise = G.generate (r*c) $ \i ->
vec `G.unsafeIndex` (offset + (i `div` c) * tda + (i `mod` c))
thaw (Matrix r c tda offset v) = MMatrix r c tda offset `liftM` G.thaw v
unsafeThaw (Matrix r c tda offset v) = MMatrix r c tda offset `liftM` G.unsafeThaw v
freeze (MMatrix r c tda offset v) = Matrix r c tda offset `liftM` G.freeze v
unsafeFreeze (MMatrix r c tda offset v) = Matrix r c tda offset `liftM` G.unsafeFreeze v
fromColumns :: G.Vector v a => [v a] -> Matrix v a
fromColumns = tr . MG.fromRows
convert :: (G.Vector v a, G.Vector w a) => Matrix v a -> Matrix w a
convert (Matrix r c tda offset vec) = Matrix r c tda offset . G.convert $ vec
subMatrix :: G.Vector v a
=> (Int, Int)
-> (Int, Int)
-> Matrix v a -> Matrix v a
subMatrix (i,j) (i',j') (Matrix _ n tda offset vec)
| m' <= 0 || n' <= 0 = MG.empty
| otherwise = Matrix m' n' tda offset' vec
where
m' = i' i + 1
n' = j' j + 1
offset' = offset + i * n + j
tr :: G.Vector v a => Matrix v a -> Matrix v a
tr (Matrix r c tda offset vec) = MG.fromVector (c,r) $ G.generate (r*c) f
where
f i = vec G.! (offset + i `mod` r * tda + i `div` r)
ident :: (Num a, G.Vector v a) => Int -> Matrix v a
ident n = diagRect 0 (n,n) $ replicate n 1
diag :: (Num a, G.Vector v a, F.Foldable t)
=> t a
-> Matrix v a
diag d = diagRect 0 (n,n) d
where n = length . F.toList $ d
diagRect :: (G.Vector v a, F.Foldable t)
=> a
-> (Int, Int)
-> t a
-> Matrix v a
diagRect z0 (r,c) d = MG.fromVector (r,c) $ G.create $ GM.replicate n z0 >>= go d c
where
go xs c' v = F.foldlM f 0 xs >> return v
where
f !i x = GM.unsafeWrite v (i*(c'+1)) x >> return (i+1)
n = r * c
fromBlocks :: G.Vector v a
=> a
-> [[Matrix v a]]
-> Matrix v a
fromBlocks d ms = MG.fromVector (m,n) $ G.create $ GM.replicate (m*n) d >>= go n ms
where
go n' xss v = foldM_ f 0 xss >> return v
where
f !cr xs = do (r', _) <- foldM g (0, 0) xs
return $ cr + r'
where
g (!maxR, !cc) x = do
let (r,c) = MG.dim x
vec = MG.flatten x
step i u = do
GM.unsafeWrite v ((cr + i `div` c) * n' + i `mod` c + cc) u
return (i+1)
G.foldM'_ step (0::Int) vec
return (max maxR r, cc + c)
(m, n) = (sum *** maximum) . unzip . Prelude.map ((maximum *** sum) .
unzip . Prelude.map (MG.rows &&& MG.cols)) $ ms
isSymmetric :: (Eq a, G.Vector v a) => Matrix v a -> Bool
isSymmetric m@(Matrix r c _ _ _) | r /= c = False
| otherwise = all f [0 .. r1]
where
f i = all g [i + 1 .. c1]
where g j = m MG.! (i,j) == m MG.! (j,i)
force :: G.Vector v a => Matrix v a -> Matrix v a
force m@(Matrix r c _ _ _) = MG.fromVector (r,c) . G.force . MG.flatten $ m
imap :: (G.Vector v a, G.Vector v b) => ((Int, Int) -> a -> b) -> Matrix v a -> Matrix v b
imap f m@(Matrix r c _ _ _) = MG.fromVector (r,c) $ G.imap f' . MG.flatten $ m
where
f' i = f (i `div` c, i `mod` c)
map :: (G.Vector v a, G.Vector v b) => (a -> b) -> Matrix v a -> Matrix v b
map f m@(Matrix r c _ _ _) = MG.fromVector (r,c) $ G.map f . MG.flatten $ m
foldl :: G.Vector v b => (a -> b -> a) -> a -> Matrix v b -> a
foldl f acc m = G.foldl f acc . MG.flatten $ m
mapM :: (G.Vector v a, G.Vector v b, Monad m) => (a -> m b) -> Matrix v a -> m (Matrix v b)
mapM f m@(Matrix r c _ _ _) = liftM (MG.fromVector (r,c)) . G.mapM f . MG.flatten $ m
mapM_ :: (G.Vector v a, Monad m) => (a -> m b) -> Matrix v a -> m ()
mapM_ f = G.mapM_ f . MG.flatten
forM :: (G.Vector v a, G.Vector v b, Monad m) => Matrix v a -> (a -> m b) -> m (Matrix v b)
forM = flip mapM
forM_ :: (G.Vector v a, Monad m) => Matrix v a -> (a -> m b) -> m ()
forM_ = flip mapM_
sequence :: (G.Vector v a, G.Vector v (m a), Monad m)
=> Matrix v (m a) -> m (Matrix v a)
sequence (Matrix r c tda offset vec) = liftM (Matrix r c tda offset) . G.sequence $ vec
sequence_ :: (G.Vector v (m a), Monad m)
=> Matrix v (m a) -> m ()
sequence_ (Matrix _ _ _ _ vec) = G.sequence_ vec
generate :: G.Vector v a => (Int, Int) -> ((Int, Int) -> a) -> Matrix v a
generate (r,c) f = MG.fromVector (r,c) . G.generate (r*c) $ \i -> f (i `div` c, i `mod` c)