module Data.Histogram.ST (
MHistogram
, newMHistogram
, fillOne
, fillOneW
, fillMonoid
, freezeHist
) where
import Control.Monad.ST
import Data.Monoid
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU
import qualified Data.Vector.Generic as G
import Data.Histogram
data MHistogram s bin a where
MHistogram :: (Bin bin, MU.Unbox a) =>
bin
-> MU.MVector s a
-> MU.MVector s a
-> MHistogram s bin a
newMHistogram :: (Bin bin, U.Unbox a) => a -> bin -> ST s (MHistogram s bin a)
newMHistogram zero bin = do
uo <- MU.newWith 2 zero
a <- MU.newWith (nBins bin) zero
return $ MHistogram bin uo a
fillOne :: Num a => MHistogram s bin a -> BinValue bin -> ST s ()
fillOne (MHistogram bin uo arr) x
| i < 0 = MU.unsafeWrite uo 0 . (+1) =<< MU.unsafeRead uo 0
| i >= MU.length arr = MU.unsafeWrite uo 1 . (+1) =<< MU.unsafeRead uo 1
| otherwise = MU.unsafeWrite arr i . (+1) =<< MU.unsafeRead arr i
where
i = toIndex bin x
fillOneW :: Num a => MHistogram s bin a -> (BinValue bin, a) -> ST s ()
fillOneW (MHistogram bin uo arr) (x,w)
| i < 0 = MU.unsafeWrite uo 0 . (+w) =<< MU.unsafeRead uo 0
| i >= MU.length arr = MU.unsafeWrite uo 1 . (+w) =<< MU.unsafeRead uo 1
| otherwise = MU.unsafeWrite arr i . (+w) =<< MU.unsafeRead arr i
where
i = toIndex bin x
fillMonoid :: Monoid a => MHistogram s bin a -> (BinValue bin, a) -> ST s ()
fillMonoid (MHistogram bin uo arr) (x,m)
| i < 0 = MU.unsafeWrite uo 1 . (flip mappend m) =<< MU.unsafeRead uo 0
| i >= MU.length arr = MU.unsafeWrite uo 1 . (flip mappend m) =<< MU.unsafeRead uo 1
| otherwise = MU.unsafeWrite arr i . (flip mappend m) =<< MU.unsafeRead arr i
where
i = toIndex bin x
freezeHist :: MHistogram s bin a -> ST s (Histogram bin a)
freezeHist (MHistogram bin uo arr) = do
u <- MU.unsafeRead uo 0
o <- MU.unsafeRead uo 1
tmp <- MU.new (MU.length arr)
MU.copy tmp arr
a <- G.unsafeFreeze tmp
return $ histogramUO bin (Just (u,o)) a