{-# LANGUAGE BangPatterns #-}
-- | Various internal utilities.
module Numeric.MCMC.Util
where
import Control.Monad.Primitive (PrimMonad, PrimState)
import System.Random.MWC
import Data.Function (fix)
import qualified Data.Sequence as Seq
import Data.Foldable (toList)
import Control.Monad
import Control.Monad.ST
import Data.STRef
import Data.List.Split (chunksOf)
-- | Given Int, bounds, and generator, generate a different Int in the bound.
genDiffInt :: PrimMonad m => Int -> (Int, Int) -> Gen (PrimState m) -> m Int
genDiffInt a bounds gen = fix $ \loopB ->
do b <- uniformR bounds gen
if a == b then loopB else return b
{-# INLINE genDiffInt #-}
-- | Tail-recursive mean function.
mean :: [Double] -> Double
mean = go 0.0 0
where go :: Double -> Int -> [Double] -> Double
go !s !l [] = s / fromIntegral l
go !s !l (x:xs) = go (s + x) (l + 1) xs
{-# INLINE mean #-}
-- | Tail-recursive mean function for lists.
listMean :: [[Double]] -> [Double]
listMean = go (repeat 0) 0
where go :: [Double] -> Int -> [[Double]] -> [Double]
go !s !l [] = map (/ fromIntegral l) s
go !s !l (xs:xss) = go (zipWith (+) s xs) (l + 1) xss
{-# INLINE listMean #-}
-- | Knuth-shuffle a list. Uses 'Seq' internally.
-- FIXME Would be nice if it was lazy, to work well with `sample`
shuffle :: [a] -> Seed -> [a]
shuffle xs seed = runST $ do
xsref <- newSTRef $ Seq.fromList xs
gen <- restore seed
let n = length xs
forM_ [n-1,n-2..1] $ \i -> do
j <- uniformR (0, i) gen
xst <- readSTRef xsref
let tmp = Seq.index xst j
modifySTRef xsref $ Seq.update j (Seq.index xst i)
modifySTRef xsref $ Seq.update i tmp
result <- readSTRef xsref
return $ toList result
{-# INLINE shuffle #-}
-- | Sample from a list without replacement.
sample :: Int -> [a] -> Seed -> [a]
sample k xs seed = take k $ shuffle xs seed
{-# INLINE sample #-}
-- | Autocovariance function. `m` is the maximum lag.
acov :: [Double] -> Int -> [Double]
acov xs m = go [] m xc
where xc = take (l - m) $ map (subtract (mean xs)) xs
l = length xs
go !s0 !j0 ys | j0 < 0 = map (/ fromIntegral (l - m)) (reverse s0)
| otherwise = go (sum (zipWith (*) xc ys) : s0) (j0 - 1) (drop 1 ys)
{-# INLINE acov #-}
-- | (mean, standard error, integrated autocorrelation time)
acor :: [Double] -> (Double, Double, Double)
acor xs = let t0 = d0 / head ac0
in if t0*fromIntegral winmult < fromIntegral maxlag
then (mean xs, sqrt (d0 / fromIntegral l0), t0)
else (mean xs, sqrt (df / fromIntegral l0), df / head acf)
where
-- Initial values and constants
(l0, ac0, d0) = (length xs, acov xs maxlag, head ac0 + 2 * sum (tail ac0))
(taumax, winmult) = (2, 5) :: (Int, Int)
maxlag = taumax*winmult
-- Final autocovariance
(sig1, acf) = go xs ac0 (sqrt (d0 / fromIntegral l0)) (d0 / head ac0)
df = 0.25 * sig1 * sig1 * fromIntegral l0
-- The recursive worker
go ys ac sig !tau
| tau*fromIntegral winmult < fromIntegral maxlag = (sig, ac)
| otherwise =
let (ys1, ac1, d, t1) = (joiner ys, acov ys1 maxlag, head ac1 + (2 :: Double) * sum (tail ac1), d / head ac1)
in if ys1 == [] then (sig, ac) else go ys1 ac1 (sqrt (d / fromIntegral (length ys))) t1
where
joiner zs = map sum $ take (truncate $ fromIntegral (length zs) / (2 :: Double)) (chunksOf 2 zs)
{-# INLINE acor #-}