{-# LANGUAGE BangPatterns, TupleSections #-}
module SMVM.Matrix (readCSRMatrix, randomCSRMatrix) where
import Random
import SMVM.MatrixMarket
import System.Random.MWC
import System.IO.Unsafe
import Data.Vector.Unboxed (Vector)
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 Int, Vector (Int,a))
-- 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 :: GenIO -> Int -> Int -> IO (CSRMatrix Float)
randomCSRMatrix gen rows cols = do
segd <- randomVectorR ( 0,cols`div`3) gen rows
let nnz = V.sum segd
inds <- randomVectorR ( 0,cols-1) gen nnz
vals <- randomVectorR (-1,1) gen nnz
return (segd, V.zip inds vals)
-- Read elements into unboxed arrays, convert to zero-indexed compress sparse
-- row format.
--
{-# INLINE csr #-}
csr :: (Int,Int) -> Int -> [(Int,Int,Float)] -> IO (Vector Int, Vector (Int,Float))
csr (m,_) l elems = do
mu <- M.new l
let 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 !n rows | n < m = let (s,ss) = V.span (==n) rows in M.unsafeWrite mseg n (V.length s) >> gos (n+1) ss
| otherwise = V.unsafeFreeze mseg
seg <- gos 0 i
return (seg , V.zip j v)
-- 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)