module Data.Histogram.ST (
HistogramST(..)
, newHistogramST
, fillOne
, fillOneW
, fillMonoid
, freezeHist
, Accumulator(..)
, Accum(Accum)
, accumList
, accumHist
, fillHistograms
) where
import Control.Monad.ST
import Data.Array.Vector
import Data.Monoid
import Data.Histogram
import Data.Histogram.Bin
data HistogramST s bin a where
HistogramST :: (Bin bin, UA a) =>
bin
-> MUArr a s
-> MUArr a s
-> HistogramST s bin a
newHistogramST :: (Bin bin, UA a) => a -> bin -> ST s (HistogramST s bin a)
newHistogramST zero bin = do
uo <- newMU 2
writeMU uo 0 zero >> writeMU uo 1 zero
a <- newMU (nBins bin)
mapM_ (\i -> writeMU a i zero) [0 .. (lengthMU a) 1]
return $ HistogramST bin uo a
fillOne :: Num a => HistogramST s bin a -> BinValue bin -> ST s ()
fillOne (HistogramST bin uo arr) x
| i < 0 = writeMU uo 0 . (+1) =<< readMU uo 0
| i >= lengthMU arr = writeMU uo 1 . (+1) =<< readMU uo 1
| otherwise = writeMU arr i . (+1) =<< readMU arr i
where
i = toIndex bin x
fillOneW :: Num a => HistogramST s bin a -> (BinValue bin, a) -> ST s ()
fillOneW (HistogramST bin uo arr) (x,w)
| i < 0 = writeMU uo 0 . (+w) =<< readMU uo 0
| i >= lengthMU arr = writeMU uo 1 . (+w) =<< readMU uo 1
| otherwise = writeMU arr i . (+w) =<< readMU arr i
where
i = toIndex bin x
fillMonoid :: Monoid a => HistogramST s bin a -> (BinValue bin, a) -> ST s ()
fillMonoid (HistogramST bin uo arr) (x,m)
| i < 0 = writeMU uo 1 . (flip mappend m) =<< readMU uo 0
| i >= lengthMU arr = writeMU uo 1 . (flip mappend m) =<< readMU uo 1
| otherwise = writeMU arr i . (flip mappend m) =<< readMU arr i
where
i = toIndex bin x
freezeHist :: HistogramST s bin a -> ST s (Histogram bin a)
freezeHist (HistogramST bin uo arr) = do
[u,o] <- fromU `fmap` unsafeFreezeAllMU uo
let len = lengthMU arr
tmp <- newMU len
memcpyOffMU arr tmp 0 0 len
a <- unsafeFreezeAllMU tmp
return $ Histogram bin (Just (u,o)) a
class Accumulator h where
putOne :: h s a b -> a -> ST s ()
extract :: Monoid b => (h s a b) -> ST s b
putMany :: Accumulator h => h s a b -> [a] -> ST s ()
putMany !h = mapM_ (putOne h)
fillHistograms :: Monoid b => (forall s . ST s (Accum s a b)) -> [a] -> b
fillHistograms h xs = runST $ do h' <- h
putMany h' xs
extract h'
data Accum s a b where
Accum :: Accumulator h => h s a b -> Accum s a b
instance Accumulator Accum where
putOne !(Accum h) !x = putOne h x
extract !(Accum h) = extract h
newtype AccumList s a b = AccumList [Accum s a b]
accumList :: [ST s (Accum s a b)] -> ST s (Accum s a b)
accumList l = (Accum . AccumList) `fmap` sequence l
instance Accumulator AccumList where
putOne !(AccumList l) !x = mapM_ (flip putOne $ x) l
extract !(AccumList l) = mconcat `fmap` mapM extract l
data AccumHist s a b where
AccumHist :: (Bin bin) =>
(a -> HistogramST s bin val -> ST s ())
-> (Histogram bin val -> b)
-> HistogramST s bin val
-> AccumHist s a b
accumHist :: (Bin bin) =>
(a -> HistogramST s bin val -> ST s ())
-> (Histogram bin val -> b)
-> HistogramST s bin val
-> ST s (Accum s a b)
accumHist inp out h = return . Accum $ AccumHist inp out h
instance Accumulator AccumHist where
putOne !(AccumHist inp _ st) !x = inp x st
extract !(AccumHist _ out st) = out `fmap` freezeHist st