{-# LANGUAGE BangPatterns, CPP, GADTs, FlexibleContexts, ScopedTypeVariables #-}
module System.Random.MWC.Distributions
(
normal
, standard
, exponential
, truncatedExp
, gamma
, chiSquare
, beta
, categorical
, logCategorical
, geometric0
, geometric1
, bernoulli
, dirichlet
, uniformPermutation
, uniformShuffle
, uniformShuffleM
) where
import Prelude hiding (mapM)
import Control.Monad (liftM)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Bits ((.&.))
import Data.Foldable (foldl')
#if !MIN_VERSION_base(4,8,0)
import Data.Traversable (Traversable)
#endif
import Data.Traversable (mapM)
import Data.Word (Word32)
import System.Random.Stateful (StatefulGen(..),Uniform(..),UniformRange(..),uniformDoublePositive01M)
import qualified Data.Vector.Unboxed as I
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as M
data T = T {-# UNPACK #-} !Double {-# UNPACK #-} !Double
normal :: StatefulGen g m
=> Double
-> Double
-> g
-> m Double
{-# INLINE normal #-}
normal m s gen = do
x <- standard gen
return $! m + s * x
standard :: StatefulGen g m => g -> m Double
{-# INLINE standard #-}
standard gen = loop
where
loop = do
u <- (subtract 1 . (*2)) `liftM` uniformDoublePositive01M gen
ri <- uniformM gen
let i = fromIntegral ((ri :: Word32) .&. 127)
bi = I.unsafeIndex blocks i
bj = I.unsafeIndex blocks (i+1)
case () of
_| abs u < I.unsafeIndex ratios i -> return $! u * bi
| i == 0 -> normalTail (u < 0)
| otherwise -> do
let x = u * bi
xx = x * x
d = exp (-0.5 * (bi * bi - xx))
e = exp (-0.5 * (bj * bj - xx))
c <- uniformDoublePositive01M gen
if e + c * (d - e) < 1
then return x
else loop
normalTail neg = tailing
where tailing = do
x <- ((/rNorm) . log) `liftM` uniformDoublePositive01M gen
y <- log `liftM` uniformDoublePositive01M gen
if y * (-2) < x * x
then tailing
else return $! if neg then x - rNorm else rNorm - x
blocks :: I.Vector Double
blocks = (`I.snoc` 0) . I.cons (v/f) . I.cons rNorm . I.unfoldrN 126 go $! T rNorm f
where
go (T b g) = let !u = T h (exp (-0.5 * h * h))
h = sqrt (-2 * log (v / b + g))
in Just (h, u)
v = 9.91256303526217e-3
f = exp (-0.5 * rNorm * rNorm)
{-# NOINLINE blocks #-}
rNorm :: Double
rNorm = 3.442619855899
ratios :: I.Vector Double
ratios = I.zipWith (/) (I.tail blocks) blocks
{-# NOINLINE ratios #-}
exponential :: StatefulGen g m
=> Double
-> g
-> m Double
{-# INLINE exponential #-}
exponential b gen = do
x <- uniformDoublePositive01M gen
return $! - log x / b
truncatedExp :: StatefulGen g m
=> Double
-> (Double,Double)
-> g
-> m Double
{-# INLINE truncatedExp #-}
truncatedExp scale (a,b) gen = do
let delta = b - a
p <- uniformDoublePositive01M gen
return $! a - log ( (1 - p) + p*exp(-scale*delta)) / scale
gamma :: (StatefulGen g m)
=> Double
-> Double
-> g
-> m Double
{-# INLINE gamma #-}
gamma a b gen
| a <= 0 = pkgError "gamma" "negative alpha parameter"
| otherwise = mainloop
where
mainloop = do
T x v <- innerloop
u <- uniformDoublePositive01M gen
let cont = u > 1 - 0.331 * sqr (sqr x)
&& log u > 0.5 * sqr x + a1 * (1 - v + log v)
case () of
_| cont -> mainloop
| a >= 1 -> return $! a1 * v * b
| otherwise -> do y <- uniformDoublePositive01M gen
return $! y ** (1 / a) * a1 * v * b
innerloop = do
x <- standard gen
case 1 + a2*x of
v | v <= 0 -> innerloop
| otherwise -> return $! T x (v*v*v)
a' = if a < 1 then a + 1 else a
a1 = a' - 1/3
a2 = 1 / sqrt(9 * a1)
chiSquare :: StatefulGen g m
=> Int
-> g
-> m Double
{-# INLINE chiSquare #-}
chiSquare n gen
| n <= 0 = pkgError "chiSquare" "number of degrees of freedom must be positive"
| otherwise = do x <- gamma (0.5 * fromIntegral n) 1 gen
return $! 2 * x
geometric0 :: StatefulGen g m
=> Double
-> g
-> m Int
{-# INLINE geometric0 #-}
geometric0 p gen
| p == 1 = return 0
| p > 0 && p < 1 = do q <- uniformDoublePositive01M gen
return $! floor $ log q / log (1 - p)
| otherwise = pkgError "geometric0" "probability out of [0,1] range"
geometric1 :: StatefulGen g m
=> Double
-> g
-> m Int
{-# INLINE geometric1 #-}
geometric1 p gen = do n <- geometric0 p gen
return $! n + 1
beta :: StatefulGen g m
=> Double
-> Double
-> g
-> m Double
{-# INLINE beta #-}
beta a b gen = do
x <- gamma a 1 gen
y <- gamma b 1 gen
return $! x / (x+y)
dirichlet :: (StatefulGen g m, Traversable t)
=> t Double
-> g
-> m (t Double)
{-# INLINE dirichlet #-}
dirichlet t gen = do
t' <- mapM (\x -> gamma x 1 gen) t
let total = foldl' (+) 0 t'
return $ fmap (/total) t'
bernoulli :: StatefulGen g m
=> Double
-> g
-> m Bool
{-# INLINE bernoulli #-}
bernoulli p gen = (<p) `liftM` uniformDoublePositive01M gen
categorical :: (StatefulGen g m, G.Vector v Double)
=> v Double
-> g
-> m Int
{-# INLINE categorical #-}
categorical v gen
| G.null v = pkgError "categorical" "empty weights!"
| otherwise = do
let cv = G.scanl1' (+) v
p <- (G.last cv *) `liftM` uniformDoublePositive01M gen
return $! case G.findIndex (>=p) cv of
Just i -> i
Nothing -> pkgError "categorical" "bad weights!"
logCategorical :: (StatefulGen g m, G.Vector v Double)
=> v Double
-> g
-> m Int
{-# INLINE logCategorical #-}
logCategorical v gen
| G.null v = pkgError "logCategorical" "empty weights!"
| otherwise = categorical (G.map (exp . subtract m) v) gen
where
m = G.maximum v
uniformPermutation :: forall g m v. (StatefulGen g m, PrimMonad m, G.Vector v Int)
=> Int
-> g
-> m (v Int)
{-# INLINE uniformPermutation #-}
uniformPermutation n gen
| n < 0 = pkgError "uniformPermutation" "size must be >=0"
| otherwise = uniformShuffle (G.generate n id :: v Int) gen
uniformShuffle :: (StatefulGen g m, PrimMonad m, G.Vector v a)
=> v a
-> g
-> m (v a)
{-# INLINE uniformShuffle #-}
uniformShuffle vec gen
| G.length vec <= 1 = return vec
| otherwise = do
mvec <- G.thaw vec
uniformShuffleM mvec gen
G.unsafeFreeze mvec
uniformShuffleM :: (StatefulGen g m, PrimMonad m, M.MVector v a)
=> v (PrimState m) a
-> g
-> m ()
{-# INLINE uniformShuffleM #-}
uniformShuffleM vec gen
| M.length vec <= 1 = return ()
| otherwise = loop 0
where
n = M.length vec
lst = n-1
loop i | i == lst = return ()
| otherwise = do j <- uniformRM (i,lst) gen
M.unsafeSwap vec i j
loop (i+1)
sqr :: Double -> Double
sqr x = x * x
{-# INLINE sqr #-}
pkgError :: String -> String -> a
pkgError func msg = error $ "System.Random.MWC.Distributions." ++ func ++
": " ++ msg