module Data.Eigen.SparseMatrix (
SparseMatrix(..),
SparseMatrixXf,
SparseMatrixXd,
SparseMatrixXcf,
SparseMatrixXcd,
values,
innerIndices,
outerStarts,
innerNNZs,
cols,
rows,
coeff,
(!),
fromList,
toList,
fromDenseList,
toDenseList,
fromMatrix,
toMatrix,
norm,
squaredNorm,
blueNorm,
block,
nonZeros,
innerSize,
outerSize,
add,
sub,
mul,
pruned,
scale,
transpose,
adjoint,
compress,
uncompress,
compressed,
encode,
decode,
thaw,
freeze,
unsafeThaw,
unsafeFreeze,
) where
import qualified Prelude as P
import Prelude hiding (map)
import qualified Data.List as L
import Data.Complex
import Foreign.C.Types
import Foreign.C.String
import Foreign.Storable
import Foreign.Ptr
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc
import Control.Monad
#if __GLASGOW_HASKELL__ >= 710
#else
import Control.Applicative
#endif
import qualified Data.Eigen.Matrix as M
import qualified Data.Eigen.Matrix.Mutable as MM
import qualified Data.Eigen.SparseMatrix.Mutable as SMM
import qualified Foreign.Concurrent as FC
import qualified Data.Eigen.Internal as I
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import qualified Data.ByteString.Internal as BSI
data SparseMatrix a b where
SparseMatrix :: I.Elem a b => !(ForeignPtr (I.CSparseMatrix a b)) -> SparseMatrix a b
type SparseMatrixXf = SparseMatrix Float CFloat
type SparseMatrixXd = SparseMatrix Double CDouble
type SparseMatrixXcf = SparseMatrix (Complex Float) (I.CComplex CFloat)
type SparseMatrixXcd = SparseMatrix (Complex Double) (I.CComplex CDouble)
instance (I.Elem a b, Show a) => Show (SparseMatrix a b) where
show m = concat [
"SparseMatrix ", show (rows m), "x", show (cols m),
"\n", L.intercalate "\n" $ P.map (L.intercalate "\t" . P.map show) $ toDenseList m, "\n"]
instance I.Elem a b => Num (SparseMatrix a b) where
(*) = mul
(+) = add
() = sub
fromInteger x = fromList 1 1 [(0,0,fromInteger x)]
signum = map signum
abs = map abs
negate = map negate
map :: I.Elem a b => (a -> a) -> SparseMatrix a b -> SparseMatrix a b
map f m = fromList (rows m) (cols m) . P.map (\(r,c,v) -> (r,c,f v)) . toList $ m
mk :: I.Elem a b => Ptr (I.CSparseMatrix a b) -> IO (SparseMatrix a b)
mk p = SparseMatrix <$> FC.newForeignPtr p (I.call $ I.sparse_free p)
values :: I.Elem a b => SparseMatrix a b -> VS.Vector b
values = _getvec I.sparse_values
innerIndices :: I.Elem a b => SparseMatrix a b -> VS.Vector CInt
innerIndices = _getvec I.sparse_innerIndices
outerStarts :: I.Elem a b => SparseMatrix a b -> VS.Vector CInt
outerStarts = _getvec I.sparse_outerStarts
innerNNZs :: I.Elem a b => SparseMatrix a b -> Maybe (VS.Vector CInt)
innerNNZs m
| compressed m = Nothing
| otherwise = Just $ _getvec I.sparse_innerNNZs m
rows :: I.Elem a b => SparseMatrix a b -> Int
rows = _unop I.sparse_rows (return . I.cast)
cols :: I.Elem a b => SparseMatrix a b -> Int
cols = _unop I.sparse_cols (return . I.cast)
coeff :: I.Elem a b => Int -> Int -> SparseMatrix a b -> a
coeff row col (SparseMatrix fp) = I.performIO $ withForeignPtr fp $ \p -> alloca $ \pq -> do
I.call $ I.sparse_coeff p (I.cast row) (I.cast col) pq
I.cast <$> peek pq
(!) :: I.Elem a b => SparseMatrix a b -> (Int, Int) -> a
(!) m (row, col) = coeff row col m
norm :: I.Elem a b => SparseMatrix a b -> a
norm = _unop I.sparse_norm (return . I.cast)
squaredNorm :: I.Elem a b => SparseMatrix a b -> a
squaredNorm = _unop I.sparse_squaredNorm (return . I.cast)
blueNorm :: I.Elem a b => SparseMatrix a b -> a
blueNorm = _unop I.sparse_blueNorm (return . I.cast)
block :: I.Elem a b => Int -> Int -> Int -> Int -> SparseMatrix a b -> SparseMatrix a b
block row col rows cols = _unop (\p pq -> I.sparse_block p (I.cast row) (I.cast col) (I.cast rows) (I.cast cols) pq) mk
nonZeros :: I.Elem a b => SparseMatrix a b -> Int
nonZeros = _unop I.sparse_nonZeros (return . I.cast)
compress :: I.Elem a b => SparseMatrix a b -> SparseMatrix a b
compress = _unop I.sparse_makeCompressed mk
uncompress :: I.Elem a b => SparseMatrix a b -> SparseMatrix a b
uncompress = _unop I.sparse_uncompress mk
compressed :: I.Elem a b => SparseMatrix a b -> Bool
compressed = _unop I.sparse_isCompressed (return . (/=0))
innerSize :: I.Elem a b => SparseMatrix a b -> Int
innerSize = _unop I.sparse_innerSize (return . I.cast)
outerSize :: I.Elem a b => SparseMatrix a b -> Int
outerSize = _unop I.sparse_outerSize (return . I.cast)
pruned :: I.Elem a b => a -> SparseMatrix a b -> SparseMatrix a b
pruned r = _unop (\p pq -> alloca $ \pr -> poke pr (I.cast r) >> I.sparse_prunedRef p pr pq) mk
scale :: I.Elem a b => a -> SparseMatrix a b -> SparseMatrix a b
scale x = _unop (\p pq -> alloca $ \px -> poke px (I.cast x) >> I.sparse_scale p px pq) mk
transpose :: I.Elem a b => SparseMatrix a b -> SparseMatrix a b
transpose = _unop I.sparse_transpose mk
adjoint :: I.Elem a b => SparseMatrix a b -> SparseMatrix a b
adjoint = _unop I.sparse_adjoint mk
add :: I.Elem a b => SparseMatrix a b -> SparseMatrix a b -> SparseMatrix a b
add = _binop I.sparse_add mk
sub :: I.Elem a b => SparseMatrix a b -> SparseMatrix a b -> SparseMatrix a b
sub = _binop I.sparse_sub mk
mul :: I.Elem a b => SparseMatrix a b -> SparseMatrix a b -> SparseMatrix a b
mul = _binop I.sparse_mul mk
fromList :: I.Elem a b => Int -> Int -> [(Int, Int, a)] -> SparseMatrix a b
fromList rows cols tris = I.performIO $ VS.unsafeWith vs $ \p -> alloca $ \pq -> do
I.call $ I.sparse_fromList (I.cast rows) (I.cast cols) p (I.cast $ VS.length vs) pq
peek pq >>= mk
where vs = VS.fromList $ P.map (\(row,col,val) -> I.CTriplet (I.cast row) (I.cast col) (I.cast val)) tris
toList :: I.Elem a b => SparseMatrix a b -> [(Int, Int, a)]
toList m@(SparseMatrix fp) = I.performIO $ do
let size = nonZeros m
tris <- VSM.new size
withForeignPtr fp $ \p ->
VSM.unsafeWith tris $ \q ->
I.call $ I.sparse_toList p q (I.cast size)
let f (I.CTriplet row col val) = (I.cast row, I.cast col, I.cast val)
P.map f . VS.toList <$> VS.unsafeFreeze tris
fromDenseList :: (I.Elem a b, Eq a) => [[a]] -> SparseMatrix a b
fromDenseList list = fromList rows cols $ do
(row, vals) <- zip [0..] list
(col, val) <- zip [0..] vals
guard $ val /= 0
return (row, col, val)
where
rows = length list
cols = L.foldl' max 0 $ P.map length list
toDenseList :: I.Elem a b => SparseMatrix a b -> [[a]]
toDenseList m = [[coeff row col m | col <- [0 .. cols m 1]] | row <- [0 .. rows m 1]]
fromMatrix :: I.Elem a b => M.Matrix a b -> SparseMatrix a b
fromMatrix m1 = I.performIO $ alloca $ \pm0 ->
M.unsafeWith m1 $ \vals rows cols -> do
I.call $ I.sparse_fromMatrix vals rows cols pm0
peek pm0 >>= mk
toMatrix :: I.Elem a b => SparseMatrix a b -> M.Matrix a b
toMatrix m1@(SparseMatrix fp) = I.performIO $ do
m0 <- MM.new (rows m1) (cols m1)
MM.unsafeWith m0 $ \vals rows cols ->
withForeignPtr fp $ \pm1 ->
I.call $ I.sparse_toMatrix pm1 vals rows cols
M.unsafeFreeze m0
encode :: I.Elem a b => SparseMatrix a b -> BSL.ByteString
encode m@(SparseMatrix fp) = I.performIO $ do
let size = nonZeros m
tris <- VSM.new size
withForeignPtr fp $ \p ->
VSM.unsafeWith tris $ \q ->
I.call $ I.sparse_toList p q (I.cast size)
tris <- VS.unsafeFreeze tris
let
tri@(I.CTriplet _ _ val) = VS.head tris
return $ BSL.fromChunks [
encodeInt (I.magicCode val),
encodeInt (I.cast $ rows m),
encodeInt (I.cast $ cols m),
encodeInt (I.cast $ size),
let (fp, fs) = VS.unsafeToForeignPtr0 tris in BSI.PS (castForeignPtr fp) 0 (fs * sizeOf tri)]
where
encodeInt :: CInt -> BS.ByteString
encodeInt x = BSI.unsafeCreate (sizeOf x) $ (`poke` x) . castPtr
decode :: forall a b . I.Elem a b => BSL.ByteString -> SparseMatrix a b
decode st = I.performIO $ do
let
val = undefined :: b
tri = undefined :: I.CTriplet b
triSize = sizeOf tri
st <- I.openStream st
code <- I.decodeInt <$> I.readStream st I.intSize
when (code /= I.magicCode val) $
fail "SparseMatrix.decode: wrong matrix type"
rows <- I.readInt st
cols <- I.readInt st
size <- I.readInt st
BSI.PS fp fo _ <- I.readStream st (I.cast size * triSize)
I.closeStream st
let tris = VS.unsafeFromForeignPtr0 (I.plusForeignPtr fp fo) (I.cast size)
VS.unsafeWith tris $ \p -> alloca $ \pq -> do
I.call $ I.sparse_fromList rows cols p size pq
peek pq >>= mk
freeze :: I.Elem a b => SMM.IOSparseMatrix a b -> IO (SparseMatrix a b)
freeze (SMM.IOSparseMatrix fp) = SparseMatrix <$> _clone fp
thaw :: I.Elem a b => SparseMatrix a b -> IO (SMM.IOSparseMatrix a b)
thaw (SparseMatrix fp) = SMM.IOSparseMatrix <$> _clone fp
unsafeFreeze :: I.Elem a b => SMM.IOSparseMatrix a b -> IO (SparseMatrix a b)
unsafeFreeze (SMM.IOSparseMatrix fp) = return $! SparseMatrix fp
unsafeThaw :: I.Elem a b => SparseMatrix a b -> IO (SMM.IOSparseMatrix a b)
unsafeThaw (SparseMatrix fp) = return $! SMM.IOSparseMatrix fp
_unop :: Storable c => (I.CSparseMatrixPtr a b -> Ptr c -> IO CString) -> (c -> IO d) -> SparseMatrix a b -> d
_unop f g (SparseMatrix fp) = I.performIO $
withForeignPtr fp $ \p ->
alloca $ \pq -> do
I.call (f p pq)
peek pq >>= g
_binop :: Storable c => (I.CSparseMatrixPtr a b -> I.CSparseMatrixPtr a b -> Ptr c -> IO CString) -> (c -> IO d) -> SparseMatrix a b -> SparseMatrix a b -> d
_binop f g (SparseMatrix fp1) (SparseMatrix fp2) = I.performIO $
withForeignPtr fp1 $ \p1 ->
withForeignPtr fp2 $ \p2 ->
alloca $ \pq -> do
I.call (f p1 p2 pq)
peek pq >>= g
_getvec :: (I.Elem a b, Storable c) => (Ptr (I.CSparseMatrix a b) -> Ptr CInt -> Ptr (Ptr c) -> IO CString) -> SparseMatrix a b -> VS.Vector c
_getvec f (SparseMatrix fm) = I.performIO $
withForeignPtr fm $ \m ->
alloca $ \ps ->
alloca $ \pq -> do
I.call $ f m ps pq
s <- fromIntegral <$> peek ps
q <- peek pq
fr <- FC.newForeignPtr q $ touchForeignPtr fm
return $! VS.unsafeFromForeignPtr0 fr s
_clone :: I.Elem a b => ForeignPtr (I.CSparseMatrix a b) -> IO (ForeignPtr (I.CSparseMatrix a b))
_clone fp = withForeignPtr fp $ \p -> alloca $ \pq -> do
I.call $ I.sparse_clone p pq
q <- peek pq
FC.newForeignPtr q $ I.call $ I.sparse_free q