module Numeric.LinearAlgebra.Sparse.IntMap where
import Numeric.LinearAlgebra.Class
import qualified Data.IntMap.Strict as IM
instance Set IM.IntMap where
liftU2 = IM.unionWith
liftI2 = IM.intersectionWith
instance Additive IM.IntMap where
zero = IM.empty
(^+^) = liftU2 (+)
instance VectorSpace IM.IntMap where
n .* im = IM.map (* n) im
instance Hilbert IM.IntMap where
a `dot` b = sum $ liftI2 (*) a b
instance Normed IM.IntMap where
norm p v | p==1 = norm1 v
| p==2 = norm2 v
| otherwise = normP p v
insertIM2 ::
IM.Key -> IM.Key -> a -> IM.IntMap (IM.IntMap a) -> IM.IntMap (IM.IntMap a)
insertIM2 i j x imm = IM.insert i ro imm where
ro = maybe (IM.singleton j x) (IM.insert j x) (IM.lookup i imm)
lookupIM2 ::
IM.Key -> IM.Key -> IM.IntMap (IM.IntMap a) -> Maybe a
lookupIM2 i j imm = IM.lookup i imm >>= IM.lookup j
fromListIM2 ::
Foldable t =>
t (IM.Key, IM.Key, a) -> IM.IntMap (IM.IntMap a) -> IM.IntMap (IM.IntMap a)
fromListIM2 iix sm = foldl ins sm iix where
ins t (i,j,x) = insertIM2 i j x t
ifoldlIM2' :: (IM.Key -> IM.Key -> a -> b -> b) -> b -> IM.IntMap (IM.IntMap a) -> b
ifoldlIM2' f empty mm = IM.foldlWithKey' accRow empty mm where
accRow acc i r = IM.foldlWithKey' (accElem i) acc r
accElem i acc j x = f i j x acc
ifoldlIM2 ::
(IM.Key -> IM.Key -> t -> IM.IntMap a -> IM.IntMap a) ->
IM.IntMap (IM.IntMap t) ->
IM.IntMap a
ifoldlIM2 f m = IM.foldlWithKey' accRow IM.empty m where
accRow acc i row = IM.foldlWithKey' (accElem i) acc row
accElem i acc j x = f i j x acc
foldlIM2 :: (a -> b -> b) -> b -> IM.IntMap (IM.IntMap a) -> b
foldlIM2 f empty mm = IM.foldl accRow empty mm where
accRow acc r = IM.foldl accElem acc r
accElem acc x = f x acc
transposeIM2 :: IM.IntMap (IM.IntMap a) -> IM.IntMap (IM.IntMap a)
transposeIM2 = ifoldlIM2 (flip insertIM2)
ifilterIM2 ::
(IM.Key -> IM.Key -> a -> Bool) ->
IM.IntMap (IM.IntMap a) ->
IM.IntMap (IM.IntMap a)
ifilterIM2 f =
IM.mapWithKey (\irow row -> IM.filterWithKey (f irow) row)
filterSubdiag :: IM.IntMap (IM.IntMap a) -> IM.IntMap (IM.IntMap a)
filterSubdiag = ifilterIM2 (\i j _ -> i>j)
countSubdiagonalNZ :: IM.IntMap (IM.IntMap a) -> Int
countSubdiagonalNZ im =
IM.size $ IM.filter (not . IM.null) (filterSubdiag im)
subdiagIndices :: IM.IntMap (IM.IntMap a) -> [(IM.Key, IM.Key)]
subdiagIndices im = concatMap rpairs $ IM.toList (IM.map IM.keys im') where
im' = filterSubdiag im
rpairs :: (a, [b]) -> [(a, b)]
rpairs (i, jj@(_:_)) = zip (replicate (length jj) i) jj
rpairs (_, []) = []
mapIM2 :: (a -> b) -> IM.IntMap (IM.IntMap a) -> IM.IntMap (IM.IntMap b)
mapIM2 = IM.map . IM.map
imapIM2 ::
(IM.Key -> IM.Key -> a -> b) ->
IM.IntMap (IM.IntMap a) ->
IM.IntMap (IM.IntMap b)
imapIM2 f im = IM.mapWithKey ff im where
ff j x = IM.mapWithKey (`f` j) x
mapKeysIM2 ::
(IM.Key -> IM.Key) -> (IM.Key -> IM.Key) -> IM.IntMap (IM.IntMap a) -> IM.IntMap (IM.IntMap a)
mapKeysIM2 fi fj im = IM.map adjCols adjRows where
adjRows = IM.mapKeys fi im
adjCols = IM.mapKeys fj
mapColumnIM2 :: (b -> b) -> IM.IntMap (IM.IntMap b) -> Int -> IM.IntMap (IM.IntMap b)
mapColumnIM2 f im jj = imapIM2 (\i j x -> if j == jj then f x else x) im