{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
module Matrix where
import Data.Int
import MatrixMarket
import System.Random.MWC
import System.IO.Unsafe
import Control.Monad.Primitive
import Data.Vector.Unboxed ( Vector, Unbox )
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Unboxed.Mutable as M
import qualified Data.Vector.Algorithms.Intro as V
type CSRMatrix a =
( Vector Int32 -- segment descriptor
, Vector (Int32, a) -- sparse vector for the (flattened) rows
, Int -- number of columns
)
-- Read a sparse matrix from a MatrixMarket file. Pattern matrices are filled
-- with random numbers in the range (-1,1).
--
{-# INLINE readCSRMatrix #-}
readCSRMatrix
:: GenIO
-> FilePath
-> IO (CSRMatrix Float)
readCSRMatrix gen file = do
mtx <- readMatrix file
case mtx of
(RealMatrix dim l vals) -> csr dim l vals
(PatternMatrix dim l ix) -> csr dim l =<< mapM' (\(a,b) -> (a,b,) `fmap` uniformR (-1,1) gen) ix
(IntMatrix _ _ _) -> error "IntMatrix type not supported"
(ComplexMatrix _ _ _) -> error "ComplexMatrix type not supported"
-- A randomly generated matrix of given size
--
{-# INLINE randomCSRMatrix #-}
randomCSRMatrix
:: (PrimMonad m, Variate a, Num a, Unbox a)
=> Gen (PrimState m)
-> Int
-> Int
-> m (CSRMatrix a)
randomCSRMatrix gen rows cols = do
segd <- randomVectorR ( 0, fromIntegral cols-1) gen rows
let nnz = fromIntegral $ V.sum segd
inds <- randomVectorR ( 0, fromIntegral cols-1) gen nnz
vals <- randomVectorR (-1,1) gen nnz
return (segd, V.zip inds vals, cols)
{-# INLINE randomVectorR #-}
randomVectorR
:: (PrimMonad m, Variate a, Unbox a)
=> (a, a)
-> Gen (PrimState m)
-> Int
-> m (Vector a)
randomVectorR r g n = V.replicateM n (uniformR r g)
-- Read elements into unboxed arrays, convert to zero-indexed compressed sparse
-- row format.
--
{-# INLINE csr #-}
csr :: forall a. (Fractional a, Unbox a)
=> (Int,Int)
-> Int
-> [(Int32,Int32,a)]
-> IO (Vector Int32, Vector (Int32,a), Int)
csr (m,_n) l elems = do
mu <- M.new l :: IO (M.IOVector (Int32,Int32,a))
let goe :: Int -> [(Int32,Int32,a)] -> IO ()
goe _ [] = return ()
goe !n (x:xs) = let (i,j,v) = x in M.unsafeWrite mu n (i-1,j-1,v) >> goe (n+1) xs
goe 0 elems
let cmp (x1,y1,_) (x2,y2,_) | x1 == x2 = compare y1 y2
| otherwise = compare x1 x2
V.sortBy cmp mu
(i,j,v) <- V.unzip3 `fmap` V.unsafeFreeze mu
mseg <- M.new m
let gos :: Int -> Vector Int32 -> IO (Vector Int32)
gos !n rows
| n >= m = V.unsafeFreeze mseg
| otherwise = let (s,ss) = V.span (== fromIntegral n) rows
in M.unsafeWrite mseg n (fromIntegral $ V.length s) >> gos (n+1) ss
seg <- gos 0 i
return (seg , V.zip j v, _n)
-- Lazier versions of things in Control.Monad
--
sequence' :: [IO a] -> IO [a]
sequence' ms = foldr k (return []) ms
where k m m' = do { x <- m; xs <- unsafeInterleaveIO m'; return (x:xs) }
mapM' :: (a -> IO b) -> [a] -> IO [b]
mapM' f as = sequence' (map f as)