{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE TypeOperators #-}
module Eigen.SparseMatrix
(
SparseMatrix(..)
, SparseMatrixXf
, SparseMatrixXd
, SparseMatrixXcf
, SparseMatrixXcd
, elems
, values
, innerIndices
, outerStarts
, innerNNZs
, cols
, rows
, coeff
, (!)
, getRow
, getCol
, fromList
, toList
, fromVector
, toVector
, fromDenseList
, toDenseList
, fromMatrix
, toMatrix
, norm
, squaredNorm
, blueNorm
, block
, nonZeros
, innerSize
, outerSize
, add
, sub
, mul
, pruned
, scale
, transpose
, adjoint
, map
, imap
, compress
, uncompress
, compressed
, encode
, decode
, thaw
, freeze
, unsafeThaw
, unsafeFreeze
) where
import Control.Monad (when, guard)
import Control.Monad.Primitive (PrimMonad(..), unsafePrimToPrim)
import qualified Prelude
import Prelude hiding (read, map)
import Data.Binary (Binary(..))
import qualified Data.Binary as Binary
import qualified Data.ByteString.Lazy as BSL
import Data.Complex (Complex)
import Data.Kind (Type)
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import Foreign.C.String (CString)
import Foreign.C.Types (CInt)
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, touchForeignPtr)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Ptr (Ptr)
import Foreign.Storable (Storable(..))
import qualified Foreign.Concurrent as FC
import GHC.TypeLits (Nat, KnownNat, type (<=))
import Eigen.Internal
( Elem
, Cast(..)
, CSparseMatrix
, CSparseMatrixPtr
, natToInt
, Row(..)
, Col(..)
, CTriplet(..)
)
import qualified Data.List as List
import qualified Eigen.Internal as Internal
import qualified Eigen.Matrix as M
import qualified Eigen.Matrix.Mutable as MM
import qualified Eigen.SparseMatrix.Mutable as SMM
newtype SparseMatrix :: Nat -> Nat -> Type -> Type where
SparseMatrix :: ForeignPtr (CSparseMatrix a) -> SparseMatrix n m a
instance forall n m a. (Elem a, Show a, KnownNat n, KnownNat m) => Show (SparseMatrix n m a) where
show m = concat
[ "SparseMatrix"
, show (rows_ m)
, "x"
, show (cols_ m)
, "\n"
, List.intercalate "\n" $ Prelude.map (List.intercalate "\t" . Prelude.map show) $ toDenseList m
, "\n"
]
instance forall n m a. (Elem a, KnownNat n, KnownNat m) => Binary (SparseMatrix n m a) where
put mat = do
put $ Internal.magicCode (undefined :: C a)
put $ natToInt @n
put $ natToInt @m
put $ toVector mat
get = do
get >>= (`when` fail "wrong matrix type") . (/= Internal.magicCode (undefined :: C a))
fromVector <$> get
encode :: (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> BSL.ByteString
encode = Binary.encode
decode :: (Elem a, KnownNat n, KnownNat m) => BSL.ByteString -> SparseMatrix n m a
decode = Binary.decode
type SparseMatrixXf n m = SparseMatrix n m Float
type SparseMatrixXd n m = SparseMatrix n m Double
type SparseMatrixXcf n m = SparseMatrix n m (Complex Float)
type SparseMatrixXcd n m = SparseMatrix n m (Complex Double)
values :: Elem a => SparseMatrix n m a -> VS.Vector a
values = VS.map fromC . _getvec Internal.sparse_values
innerIndices :: Elem a => SparseMatrix n m a -> VS.Vector Int
innerIndices = VS.map fromC . _getvec Internal.sparse_innerIndices
outerStarts :: Elem a => SparseMatrix n m a -> VS.Vector Int
outerStarts = VS.map fromC . _getvec Internal.sparse_outerStarts
innerNNZs :: Elem a => SparseMatrix n m a -> Maybe (VS.Vector Int)
innerNNZs m
| compressed m = Nothing
| otherwise = Just $ VS.map fromC $ _getvec Internal.sparse_innerNNZs m
rows :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Row n
rows _ = Row @n
cols :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Col m
cols _ = Col @m
rows_ :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Int
rows_ _ = natToInt @n
cols_ :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Int
cols_ _ = natToInt @m
coeff :: forall n m r c a. (Elem a, KnownNat n, KnownNat m, KnownNat r, KnownNat c, r <= n, c <= m)
=> Row r -> Col c -> SparseMatrix n m a -> a
coeff _ _ (SparseMatrix fp) =
let !c_row = toC $! natToInt @r
!c_col = toC $! natToInt @c
in Internal.performIO $ withForeignPtr fp $ \p -> alloca $ \pq -> do
Internal.call $ Internal.sparse_coeff p c_row c_col pq
fromC <$> peek pq
(!) :: forall n m r c a. (Elem a, KnownNat n, KnownNat m, KnownNat r, KnownNat c, r <= n, c <= m)
=> SparseMatrix n m a -> (Row r, Col c) -> a
(!) m (row,col) = coeff row col m
norm :: Elem a => SparseMatrix n m a -> a
norm = _unop Internal.sparse_norm (pure . fromC)
squaredNorm :: Elem a => SparseMatrix n m a -> a
squaredNorm = _unop Internal.sparse_squaredNorm (pure . fromC)
blueNorm :: Elem a => SparseMatrix n m a -> a
blueNorm = _unop Internal.sparse_blueNorm (pure . fromC)
block :: forall sr sc br bc n m a.
(Elem a, KnownNat sr, KnownNat sc, KnownNat br, KnownNat bc, KnownNat n, KnownNat m)
=> (sr <= n, sc <= m, br <= n, bc <= m)
=> Row sr
-> Col sc
-> Row br
-> Col bc
-> SparseMatrix n m a
-> SparseMatrix br bc a
block _ _ _ _ =
let !c_startRow = toC $! natToInt @sr
!c_startCol = toC $! natToInt @sc
!c_rows = toC $! natToInt @br
!c_cols = toC $! natToInt @bc
in _unop (\p pq -> Internal.sparse_block p c_startRow c_startCol c_rows c_cols pq) _mk
nonZeros :: Elem a => SparseMatrix n m a -> Int
nonZeros = _unop Internal.sparse_nonZeros (pure . fromC)
elems :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Int
elems _ = (natToInt @n) * (natToInt @m)
compress :: Elem a => SparseMatrix n m a -> SparseMatrix n m a
compress = _unop Internal.sparse_makeCompressed _mk
uncompress :: Elem a => SparseMatrix n m a -> SparseMatrix n m a
uncompress = _unop Internal.sparse_uncompress _mk
compressed :: Elem a => SparseMatrix n m a -> Bool
compressed = _unop Internal.sparse_isCompressed (pure . (/=0))
innerSize :: Elem a => SparseMatrix n m a -> Int
innerSize = _unop Internal.sparse_innerSize (pure . fromC)
outerSize :: Elem a => SparseMatrix n m a -> Int
outerSize = _unop Internal.sparse_outerSize (pure . fromC)
pruned :: Elem a => a -> SparseMatrix n m a -> SparseMatrix n m a
pruned r = _unop (\p pq -> alloca $ \pr -> poke pr (toC r) >> Internal.sparse_prunedRef p pr pq) _mk
scale :: Elem a => a -> SparseMatrix n m a -> SparseMatrix n m a
scale x = _unop (\p pq -> alloca $ \px -> poke px (toC x) >> Internal.sparse_scale p px pq) _mk
transpose :: Elem a => SparseMatrix n m a -> SparseMatrix m n a
transpose = _unop Internal.sparse_transpose _mk
adjoint :: Elem a => SparseMatrix n m a -> SparseMatrix m n a
adjoint = _unop Internal.sparse_adjoint _mk
add :: Elem a => SparseMatrix n m a -> SparseMatrix n m a -> SparseMatrix n m a
add = _binop Internal.sparse_add _mk
sub :: Elem a => SparseMatrix n m a -> SparseMatrix n m a -> SparseMatrix n m a
sub = _binop Internal.sparse_sub _mk
mul :: Elem a => SparseMatrix p q a -> SparseMatrix q r a -> SparseMatrix p r a
mul = _binop Internal.sparse_mul _mk
map :: (Elem a, Elem b, KnownNat n, KnownNat m) => (a -> b) -> SparseMatrix n m a -> SparseMatrix n m b
map f m = fromVector . VS.map g . toVector $ m where
g (CTriplet r c v) = CTriplet r c $ (toC . f . fromC) v
imap :: (Elem a, Elem b, KnownNat n, KnownNat m) => (Int -> Int -> a -> b) -> SparseMatrix n m a -> SparseMatrix n m b
imap f m = fromVector . VS.map g . toVector $ m where
g (CTriplet r c v) =
let !_r = fromC r
!_c = fromC c
!_v = fromC v
in CTriplet r c $ toC $ f _r _c _v
fromVector :: forall n m a. (Elem a, KnownNat n, KnownNat m)
=> VS.Vector (CTriplet a)
-> SparseMatrix n m a
fromVector tris =
let !c_rs = toC $! natToInt @n
!c_cs = toC $! natToInt @m
!len = toC $! VS.length tris
in Internal.performIO $ VS.unsafeWith tris $ \p ->
alloca $ \pq -> do
Internal.call $ Internal.sparse_fromList c_rs c_cs p len pq
peek pq >>= _mk
toVector :: Elem a => SparseMatrix n m a -> VS.Vector (CTriplet a)
toVector m@(SparseMatrix fp) = Internal.performIO $ do
let !size = nonZeros m
tris <- VSM.new size
withForeignPtr fp $ \p ->
VSM.unsafeWith tris $ \q ->
Internal.call $ Internal.sparse_toList p q (toC size)
VS.unsafeFreeze tris
toList :: Elem a => SparseMatrix n m a -> [(Int, Int, a)]
toList = Prelude.map fromC . VS.toList . toVector
fromList :: (Elem a, KnownNat n, KnownNat m) => [(Int, Int, a)] -> SparseMatrix n m a
fromList = fromVector . VS.fromList . fmap toC
toDenseList :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> [[a]]
toDenseList mat = [[_unsafeCoeff row col mat | col <- [0 .. _unsafeCols mat - 1]] | row <- [0 .. _unsafeRows mat - 1]]
fromDenseList :: forall n m a. (Elem a, Eq a, KnownNat n, KnownNat m) => [[a]] -> Maybe (SparseMatrix n m a)
fromDenseList list =
let _rows = List.length list
_cols = List.foldl' max 0 $ List.map length list
in if ((_rows /= (natToInt @n)) || (_cols /= (natToInt @m)))
then Nothing
else Just $ fromList $ do
(row, vals) <- zip [0..] list
(col, val) <- zip [0..] vals
guard $ val /= 0
return (row, col, val)
toMatrix :: (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> M.Matrix n m a
toMatrix (SparseMatrix fp) = Internal.performIO $ do
m0 :: MM.IOMatrix n m a <- MM.new
MM.unsafeWith m0 $ \_vals _rows _cols ->
withForeignPtr fp $ \pm1 ->
Internal.call $ Internal.sparse_toMatrix pm1 _vals _rows _cols
M.unsafeFreeze m0
fromMatrix :: (Elem a, KnownNat n, KnownNat m) => M.Matrix n m a -> SparseMatrix n m a
fromMatrix m1 = Internal.performIO $ alloca $ \pm0 ->
M.unsafeWith m1 $ \_vals _rows _cols -> do
Internal.call $ Internal.sparse_fromMatrix _vals _rows _cols pm0
peek pm0 >>= _mk
freeze :: (Elem a, PrimMonad p) => SMM.MSparseMatrix n m (PrimState p) a -> p (SparseMatrix n m a)
freeze (SMM.MSparseMatrix fp) = SparseMatrix <$> _clone fp
thaw :: (Elem a, PrimMonad p) => SparseMatrix n m a -> p (SMM.MSparseMatrix n m (PrimState p) a)
thaw (SparseMatrix fp) = SMM.MSparseMatrix <$> _clone fp
unsafeFreeze :: (Elem a, PrimMonad p) => SMM.MSparseMatrix n m (PrimState p) a -> p (SparseMatrix n m a)
unsafeFreeze (SMM.MSparseMatrix fp) = return $! SparseMatrix fp
unsafeThaw :: (Elem a, PrimMonad p) => SparseMatrix n m a -> p (SMM.MSparseMatrix n m (PrimState p) a)
unsafeThaw (SparseMatrix fp) = return $! SMM.MSparseMatrix fp
getRow :: forall n m r a. (Elem a, KnownNat n, KnownNat m, KnownNat r, r <= n, 1 <= n) => Row r -> SparseMatrix n m a -> SparseMatrix 1 m a
getRow row mat = block row (Col @0) (Row @1) (Col @m) mat
getCol :: forall n m c a. (Elem a, KnownNat n, KnownNat m, KnownNat c, c <= m, 1 <= m) => Col c -> SparseMatrix n m a -> SparseMatrix n 1 a
getCol col mat = block (Row @0) col (Row @n) (Col @1) mat
_unop :: Storable b => (CSparseMatrixPtr a -> Ptr b -> IO CString) -> (b -> IO c) -> SparseMatrix n m a -> c
_unop f g (SparseMatrix fp) = Internal.performIO $
withForeignPtr fp $ \p ->
alloca $ \pq -> do
Internal.call (f p pq)
peek pq >>= g
_binop :: Storable b => (CSparseMatrixPtr a -> CSparseMatrixPtr a -> Ptr b -> IO CString) -> (b -> IO c) -> SparseMatrix n m a -> SparseMatrix n1 m1 a -> c
_binop f g (SparseMatrix fp1) (SparseMatrix fp2) = Internal.performIO $
withForeignPtr fp1 $ \p1 ->
withForeignPtr fp2 $ \p2 ->
alloca $ \pq -> do
Internal.call (f p1 p2 pq)
peek pq >>= g
_getvec :: (Elem a, Storable b) => (Ptr (CSparseMatrix a) -> Ptr CInt -> Ptr (Ptr b) -> IO CString) -> SparseMatrix n m a -> VS.Vector b
_getvec f (SparseMatrix fm) = Internal.performIO $
withForeignPtr fm $ \m ->
alloca $ \ps ->
alloca $ \pq -> do
Internal.call $ f m ps pq
s <- fromIntegral <$> peek ps
q <- peek pq
fr <- FC.newForeignPtr q $ touchForeignPtr fm
pure $! VS.unsafeFromForeignPtr0 fr s
_clone :: (PrimMonad p, Elem a) => ForeignPtr (CSparseMatrix a) -> p (ForeignPtr (CSparseMatrix a))
_clone fp = unsafePrimToPrim $ withForeignPtr fp $ \p -> alloca $ \pq -> do
Internal.call $ Internal.sparse_clone p pq
q <- peek pq
FC.newForeignPtr q $ Internal.call $ Internal.sparse_free q
_mk :: Elem a => Ptr (CSparseMatrix a) -> IO (SparseMatrix n m a)
_mk p = SparseMatrix <$> FC.newForeignPtr p (Internal.call $ Internal.sparse_free p)
_unsafeRows :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Int
{-# INLINE _unsafeRows #-}
_unsafeRows _ = natToInt @n
_unsafeCols :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Int
{-# INLINE _unsafeCols #-}
_unsafeCols _ = natToInt @m
_unsafeCoeff :: (Elem a, KnownNat n) => Int -> Int -> SparseMatrix n m a -> a
{-# INLINE _unsafeCoeff #-}
_unsafeCoeff !row !col (SparseMatrix fp) =
let !c_row = toC row
!c_col = toC col
in Internal.performIO $ withForeignPtr fp $ \p -> alloca $ \pq -> do
Internal.call $ Internal.sparse_coeff p c_row c_col pq
fromC <$> peek pq