{-# LANGUAGE GADTs #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE Rank2Types #-}
-- |
-- Module : Data.Histogram.ST
-- Copyright : Copyright (c) 2009, Alexey Khudyakov <alexey.skladnoy@gmail.com>
-- License : BSD3
-- Maintainer : Alexey Khudyakov <alexey.skladnoy@gmail.com>
-- Stability : experimental
--
-- Mutable histograms.
module Data.Histogram.ST ( -- * Mutable histograms
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
----------------------------------------------------------------
-- Mutable histograms
----------------------------------------------------------------
-- | Mutable histogram.
data MHistogram s bin a where
MHistogram :: (Bin bin, MU.Unbox a) =>
bin -- Binning
-> MU.MVector s a -- Over/underflows
-> MU.MVector s a -- Data
-> MHistogram s bin a
-- | Create new mutable histogram. All bins are set to zero element as
-- passed to function.
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
{-# INLINE newMHistogram #-}
-- | Put one value into histogram
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
{-# INLINE fillOne #-}
-- | Put one value into histogram with weight
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
{-# INLINE fillOneW #-}
-- | Put one monoidal element
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
{-# fillMonoid #-}
-- | Create immutable histogram from mutable one. This operation involve copying.
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
-- Copy array
tmp <- MU.new (MU.length arr)
MU.copy tmp arr
a <- G.unsafeFreeze tmp
return $ histogramUO bin (Just (u,o)) a
{-# INLINE freezeHist #-}