{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE BangPatterns, CPP, GADTs, FlexibleContexts, ScopedTypeVariables #-}
-- |
-- Module    : System.Random.MWC.Distributions
-- Copyright : (c) 2012 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Pseudo-random number generation for non-uniform distributions.

module System.Random.MWC.Distributions
    (
    -- * Variates: non-uniformly distributed values
    -- ** Continuous distributions
      normal
    , standard
    , exponential
    , truncatedExp
    , gamma
    , chiSquare
    , beta
      -- ** Discrete distribution
    , categorical
    , logCategorical
    , geometric0
    , geometric1
    , bernoulli
    , binomial
      -- ** Multivariate
    , dirichlet
      -- * Permutations
    , uniformPermutation
    , uniformShuffle
    , uniformShuffleM
    -- * References
    -- $references
    ) 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
import Numeric.SpecFunctions (logFactorial)

-- Unboxed 2-tuple
data T = T {-# UNPACK #-} !Double {-# UNPACK #-} !Double


-- | Generate a normally distributed random variate with given mean
-- and standard deviation.
normal :: StatefulGen g m
       => Double                -- ^ Mean
       -> Double                -- ^ Standard deviation
       -> g
       -> m Double
{-# INLINE normal #-}
normal :: forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
normal Double
m Double
s g
gen = do
  Double
x <- g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
standard g
gen
  Double -> m Double
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> m Double) -> Double -> m Double
forall a b. (a -> b) -> a -> b
$! Double
m Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x

-- | Generate a normally distributed random variate with zero mean and
-- unit variance.
--
-- The implementation uses Doornik's modified ziggurat algorithm.
-- Compared to the ziggurat algorithm usually used, this is slower,
-- but generates more independent variates that pass stringent tests
-- of randomness.
standard :: StatefulGen g m => g -> m Double
{-# INLINE standard #-}
standard :: forall g (m :: * -> *). StatefulGen g m => g -> m Double
standard g
gen = m Double
loop
  where
    loop :: m Double
loop = do
      Double
u  <- (Double -> Double -> Double
forall a. Num a => a -> a -> a
subtract Double
1 (Double -> Double) -> (Double -> Double) -> Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
2)) (Double -> Double) -> m Double -> m Double
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
      Word32
ri <- g -> m Word32
forall a g (m :: * -> *). (Uniform a, StatefulGen g m) => g -> m a
forall g (m :: * -> *). StatefulGen g m => g -> m Word32
uniformM g
gen
      let i :: Int
i  = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word32
ri :: Word32) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. Word32
127)
          bi :: Double
bi = Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
I.unsafeIndex Vector Double
blocks Int
i
          bj :: Double
bj = Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
I.unsafeIndex Vector Double
blocks (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
      case () of
        ()
_| Double -> Double
forall a. Num a => a -> a
abs Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
I.unsafeIndex Vector Double
ratios Int
i -> Double -> m Double
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> m Double) -> Double -> m Double
forall a b. (a -> b) -> a -> b
$! Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
bi
         | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0                         -> Bool -> m Double
normalTail (Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0)
         | Bool
otherwise                      -> do
             let x :: Double
x  = Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
bi
                 xx :: Double
xx = Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x
                 d :: Double
d  = Double -> Double
forall a. Floating a => a -> a
exp (-Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
bi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
bi Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
xx))
                 e :: Double
e  = Double -> Double
forall a. Floating a => a -> a
exp (-Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
bj Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
bj Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
xx))
             Double
c <- g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
             if Double
e Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
d Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
e) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1
               then Double -> m Double
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Double
x
               else m Double
loop
    normalTail :: Bool -> m Double
normalTail Bool
neg  = m Double
tailing
      where tailing :: m Double
tailing  = do
              Double
x <- ((Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
rNorm) (Double -> Double) -> (Double -> Double) -> Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double
forall a. Floating a => a -> a
log) (Double -> Double) -> m Double -> m Double
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
              Double
y <- Double -> Double
forall a. Floating a => a -> a
log              (Double -> Double) -> m Double -> m Double
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
              if Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
* (-Double
2) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x
                then m Double
tailing
                else Double -> m Double
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> m Double) -> Double -> m Double
forall a b. (a -> b) -> a -> b
$! if Bool
neg then Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
rNorm else Double
rNorm Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
x

-- Constants used by standard/normal. They are floated to the top
-- level to avoid performance regression (Bug #16) when blocks/ratios
-- are recalculated on each call to standard/normal. It's also
-- somewhat difficult to trigger reliably.
blocks :: I.Vector Double
blocks :: Vector Double
blocks = (Vector Double -> Double -> Vector Double
forall a. Unbox a => Vector a -> a -> Vector a
`I.snoc` Double
0) (Vector Double -> Vector Double)
-> (T -> Vector Double) -> T -> Vector Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Vector Double -> Vector Double
forall a. Unbox a => a -> Vector a -> Vector a
I.cons (Double
vDouble -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
f) (Vector Double -> Vector Double)
-> (T -> Vector Double) -> T -> Vector Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Vector Double -> Vector Double
forall a. Unbox a => a -> Vector a -> Vector a
I.cons Double
rNorm (Vector Double -> Vector Double)
-> (T -> Vector Double) -> T -> Vector Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (T -> Maybe (Double, T)) -> T -> Vector Double
forall a b. Unbox a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
I.unfoldrN Int
126 T -> Maybe (Double, T)
go (T -> Vector Double) -> T -> Vector Double
forall a b. (a -> b) -> a -> b
$! Double -> Double -> T
T Double
rNorm Double
f
  where
    go :: T -> Maybe (Double, T)
go (T Double
b Double
g) = let !u :: T
u = Double -> Double -> T
T Double
h (Double -> Double
forall a. Floating a => a -> a
exp (-Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h))
                     h :: Double
h  = Double -> Double
forall a. Floating a => a -> a
sqrt (-Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log (Double
v Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
g))
                 in (Double, T) -> Maybe (Double, T)
forall a. a -> Maybe a
Just (Double
h, T
u)
    v :: Double
v = Double
9.91256303526217e-3
    f :: Double
f = Double -> Double
forall a. Floating a => a -> a
exp (-Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
rNorm Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
rNorm)
{-# NOINLINE blocks #-}

rNorm :: Double
rNorm :: Double
rNorm = Double
3.442619855899

ratios :: I.Vector Double
ratios :: Vector Double
ratios = (Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
I.zipWith Double -> Double -> Double
forall a. Fractional a => a -> a -> a
(/) (Vector Double -> Vector Double
forall a. Unbox a => Vector a -> Vector a
I.tail Vector Double
blocks) Vector Double
blocks
{-# NOINLINE ratios #-}



-- | Generate an exponentially distributed random variate.
exponential :: StatefulGen g m
            => Double            -- ^ Scale parameter
            -> g                 -- ^ Generator
            -> m Double
{-# INLINE exponential #-}
exponential :: forall g (m :: * -> *). StatefulGen g m => Double -> g -> m Double
exponential Double
b g
gen = do
  Double
x <- g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
  Double -> m Double
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> m Double) -> Double -> m Double
forall a b. (a -> b) -> a -> b
$! - Double -> Double
forall a. Floating a => a -> a
log Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
b


-- | Generate truncated exponentially distributed random variate.
truncatedExp :: StatefulGen g m
             => Double            -- ^ Scale parameter
             -> (Double,Double)   -- ^ Range to which distribution is
                                  --   truncated. Values may be negative.
             -> g                 -- ^ Generator.
             -> m Double
{-# INLINE truncatedExp #-}
truncatedExp :: forall g (m :: * -> *).
StatefulGen g m =>
Double -> (Double, Double) -> g -> m Double
truncatedExp Double
scale (Double
a,Double
b) g
gen = do
  -- We shift a to 0 and then generate distribution truncated to [0,b-a]
  -- It's easier
  let delta :: Double
delta = Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a
  Double
p <- g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
  Double -> m Double
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> m Double) -> Double -> m Double
forall a b. (a -> b) -> a -> b
$! Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
log ( (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
pDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double -> Double
forall a. Floating a => a -> a
exp(-Double
scaleDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double
delta)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
scale

-- | Random variate generator for gamma distribution.
gamma :: (StatefulGen g m)
      => Double                 -- ^ Shape parameter
      -> Double                 -- ^ Scale parameter
      -> g                      -- ^ Generator
      -> m Double
{-# INLINE gamma #-}
gamma :: forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
gamma Double
a Double
b g
gen
  | Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0    = String -> String -> m Double
forall a. String -> String -> a
pkgError String
"gamma" String
"negative alpha parameter"
  | Bool
otherwise = m Double
mainloop
    where
      mainloop :: m Double
mainloop = do
        T Double
x Double
v <- m T
innerloop
        Double
u     <- g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
        let cont :: Bool
cont =  Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
0.331 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
sqr (Double -> Double
sqr Double
x)
                 Bool -> Bool -> Bool
&& Double -> Double
forall a. Floating a => a -> a
log Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
sqr Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
a1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log Double
v) -- Rarely evaluated
        case () of
          ()
_| Bool
cont      -> m Double
mainloop
           | Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
1    -> Double -> m Double
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> m Double) -> Double -> m Double
forall a b. (a -> b) -> a -> b
$! Double
a1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b
           | Bool
otherwise -> do Double
y <- g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
                             Double -> m Double
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> m Double) -> Double -> m Double
forall a b. (a -> b) -> a -> b
$! Double
y Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
a) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b
      -- inner loop
      innerloop :: m T
innerloop = do
        Double
x <- g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
standard g
gen
        case Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
a2Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
x of
          Double
v | Double
v Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0    -> m T
innerloop
            | Bool
otherwise -> T -> m T
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (T -> m T) -> T -> m T
forall a b. (a -> b) -> a -> b
$! Double -> Double -> T
T Double
x (Double
vDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double
vDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double
v)
      -- constants
      a' :: Double
a' = if Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1 then Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1 else Double
a
      a1 :: Double
a1 = Double
a' Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
3
      a2 :: Double
a2 = Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt(Double
9 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a1)


-- | Random variate generator for the chi square distribution.
chiSquare :: StatefulGen g m
          => Int                -- ^ Number of degrees of freedom
          -> g                  -- ^ Generator
          -> m Double
{-# INLINE chiSquare #-}
chiSquare :: forall g (m :: * -> *). StatefulGen g m => Int -> g -> m Double
chiSquare Int
n g
gen
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0    = String -> String -> m Double
forall a. String -> String -> a
pkgError String
"chiSquare" String
"number of degrees of freedom must be positive"
  | Bool
otherwise = do Double
x <- Double -> Double -> g -> m Double
forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
gamma (Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) Double
1 g
gen
                   Double -> m Double
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> m Double) -> Double -> m Double
forall a b. (a -> b) -> a -> b
$! Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x

-- | Random variate generator for the geometric distribution,
-- computing the number of failures before success. Distribution's
-- support is [0..].
geometric0 :: StatefulGen g m
           => Double            -- ^ /p/ success probability lies in (0,1]
           -> g                 -- ^ Generator
           -> m Int
{-# INLINE geometric0 #-}
geometric0 :: forall g (m :: * -> *). StatefulGen g m => Double -> g -> m Int
geometric0 Double
p g
gen
  | Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1          = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0
  | Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>  Double
0 Bool -> Bool -> Bool
&& Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1 = do Double
q <- g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
                         -- FIXME: We want to use log1p here but it will
                         --        introduce dependency on math-functions.
                         Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$! Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Double -> Double
forall a. Floating a => a -> a
log Double
q Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
log (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p)
  | Bool
otherwise       = String -> String -> m Int
forall a. String -> String -> a
pkgError String
"geometric0" String
"probability out of [0,1] range"

-- | Random variate generator for geometric distribution for number of
-- trials. Distribution's support is [1..] (i.e. just 'geometric0'
-- shifted by 1).
geometric1 :: StatefulGen g m
           => Double            -- ^ /p/ success probability lies in (0,1]
           -> g                 -- ^ Generator
           -> m Int
{-# INLINE geometric1 #-}
geometric1 :: forall g (m :: * -> *). StatefulGen g m => Double -> g -> m Int
geometric1 Double
p g
gen = do Int
n <- Double -> g -> m Int
forall g (m :: * -> *). StatefulGen g m => Double -> g -> m Int
geometric0 Double
p g
gen
                      Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$! Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

-- | Random variate generator for Beta distribution
beta :: StatefulGen g m
     => Double            -- ^ alpha (>0)
     -> Double            -- ^ beta  (>0)
     -> g                 -- ^ Generator
     -> m Double
{-# INLINE beta #-}
beta :: forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
beta Double
a Double
b g
gen = do
  Double
x <- Double -> Double -> g -> m Double
forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
gamma Double
a Double
1 g
gen
  Double
y <- Double -> Double -> g -> m Double
forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
gamma Double
b Double
1 g
gen
  Double -> m Double
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> m Double) -> Double -> m Double
forall a b. (a -> b) -> a -> b
$! Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
xDouble -> Double -> Double
forall a. Num a => a -> a -> a
+Double
y)

-- | Random variate generator for Dirichlet distribution
dirichlet :: (StatefulGen g m, Traversable t)
          => t Double          -- ^ container of parameters
          -> g                 -- ^ Generator
          -> m (t Double)
{-# INLINE dirichlet #-}
dirichlet :: forall g (m :: * -> *) (t :: * -> *).
(StatefulGen g m, Traversable t) =>
t Double -> g -> m (t Double)
dirichlet t Double
t g
gen = do
  t Double
t' <- (Double -> m Double) -> t Double -> m (t Double)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> t a -> m (t b)
mapM (\Double
x -> Double -> Double -> g -> m Double
forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
gamma Double
x Double
1 g
gen) t Double
t
  let total :: Double
total = (Double -> Double -> Double) -> Double -> t Double -> Double
forall b a. (b -> a -> b) -> b -> t a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Double
0 t Double
t'
  t Double -> m (t Double)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (t Double -> m (t Double)) -> t Double -> m (t Double)
forall a b. (a -> b) -> a -> b
$ (Double -> Double) -> t Double -> t Double
forall a b. (a -> b) -> t a -> t b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
total) t Double
t'

-- | Random variate generator for Bernoulli distribution
bernoulli :: StatefulGen g m
          => Double            -- ^ Probability of success (returning True)
          -> g                 -- ^ Generator
          -> m Bool
{-# INLINE bernoulli #-}
bernoulli :: forall g (m :: * -> *). StatefulGen g m => Double -> g -> m Bool
bernoulli Double
p g
gen = (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<Double
p) (Double -> Bool) -> m Double -> m Bool
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen

-- | Random variate generator for categorical distribution.
--
--   Note that if you need to generate a lot of variates functions
--   "System.Random.MWC.CondensedTable" will offer better
--   performance.  If only few is needed this function will faster
--   since it avoids costs of setting up table.
categorical :: (StatefulGen g m, G.Vector v Double)
            => v Double          -- ^ List of weights [>0]
            -> g                 -- ^ Generator
            -> m Int
{-# INLINE categorical #-}
categorical :: forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical v Double
v g
gen
    | v Double -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v Double
v = String -> String -> m Int
forall a. String -> String -> a
pkgError String
"categorical" String
"empty weights!"
    | Bool
otherwise = do
        let cv :: v Double
cv  = (Double -> Double -> Double) -> v Double -> v Double
forall (v :: * -> *) a. Vector v a => (a -> a -> a) -> v a -> v a
G.scanl1' Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) v Double
v
        Double
p <- (v Double -> Double
forall (v :: * -> *) a. Vector v a => v a -> a
G.last v Double
cv Double -> Double -> Double
forall a. Num a => a -> a -> a
*) (Double -> Double) -> m Double -> m Double
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
        Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$! case (Double -> Bool) -> v Double -> Maybe Int
forall (v :: * -> *) a.
Vector v a =>
(a -> Bool) -> v a -> Maybe Int
G.findIndex (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>=Double
p) v Double
cv of
                    Just Int
i  -> Int
i
                    Maybe Int
Nothing -> String -> String -> Int
forall a. String -> String -> a
pkgError String
"categorical" String
"bad weights!"

-- | Random variate generator for categorical distribution where the
--   weights are in the log domain. It's implemented in terms of
--   'categorical'.
logCategorical :: (StatefulGen g m, G.Vector v Double)
               => v Double          -- ^ List of logarithms of weights
               -> g                 -- ^ Generator
               -> m Int
{-# INLINE logCategorical #-}
logCategorical :: forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
logCategorical v Double
v g
gen
  | v Double -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v Double
v  = String -> String -> m Int
forall a. String -> String -> a
pkgError String
"logCategorical" String
"empty weights!"
  | Bool
otherwise = v Double -> g -> m Int
forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical ((Double -> Double) -> v Double -> v Double
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map (Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double) -> (Double -> Double) -> Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double -> Double
forall a. Num a => a -> a -> a
subtract Double
m) v Double
v) g
gen
  where
    m :: Double
m = v Double -> Double
forall (v :: * -> *) a. (Vector v a, Ord a) => v a -> a
G.maximum v Double
v

-- | Random variate generator for uniformly distributed permutations.
--   It returns random permutation of vector /[0 .. n-1]/.
--
--   This is the Fisher-Yates shuffle
uniformPermutation :: forall g m v. (StatefulGen g m, PrimMonad m, G.Vector v Int)
                   => Int
                   -> g
                   -> m (v Int)
{-# INLINE uniformPermutation #-}
uniformPermutation :: forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, PrimMonad m, Vector v Int) =>
Int -> g -> m (v Int)
uniformPermutation Int
n g
gen
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0     = String -> String -> m (v Int)
forall a. String -> String -> a
pkgError String
"uniformPermutation" String
"size must be >=0"
  | Bool
otherwise = v Int -> g -> m (v Int)
forall g (m :: * -> *) (v :: * -> *) a.
(StatefulGen g m, PrimMonad m, Vector v a) =>
v a -> g -> m (v a)
uniformShuffle (Int -> (Int -> Int) -> v Int
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
G.generate Int
n Int -> Int
forall a. a -> a
id :: v Int) g
gen

-- | Random variate generator for a uniformly distributed shuffle (all
--   shuffles are equiprobable) of a vector. It uses Fisher-Yates
--   shuffle algorithm.
uniformShuffle :: (StatefulGen g m, PrimMonad m, G.Vector v a)
               => v a
               -> g
               -> m (v a)
{-# INLINE uniformShuffle #-}
uniformShuffle :: forall g (m :: * -> *) (v :: * -> *) a.
(StatefulGen g m, PrimMonad m, Vector v a) =>
v a -> g -> m (v a)
uniformShuffle v a
vec g
gen
  | v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
vec Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = v a -> m (v a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return v a
vec
  | Bool
otherwise         = do
      Mutable v (PrimState m) a
mvec <- v a -> m (Mutable v (PrimState m) a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
G.thaw v a
vec
      Mutable v (PrimState m) a -> g -> m ()
forall g (m :: * -> *) (v :: * -> * -> *) a.
(StatefulGen g m, PrimMonad m, MVector v a) =>
v (PrimState m) a -> g -> m ()
uniformShuffleM Mutable v (PrimState m) a
mvec g
gen
      Mutable v (PrimState m) a -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze Mutable v (PrimState m) a
mvec

-- | In-place uniformly distributed shuffle (all shuffles are
--   equiprobable)of a vector.
uniformShuffleM :: (StatefulGen g m, PrimMonad m, M.MVector v a)
                => v (PrimState m) a
                -> g
                -> m ()
{-# INLINE uniformShuffleM #-}
uniformShuffleM :: forall g (m :: * -> *) (v :: * -> * -> *) a.
(StatefulGen g m, PrimMonad m, MVector v a) =>
v (PrimState m) a -> g -> m ()
uniformShuffleM v (PrimState m) a
vec g
gen
  | v (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
M.length v (PrimState m) a
vec Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  | Bool
otherwise         = Int -> m ()
loop Int
0
  where
    n :: Int
n   = v (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
M.length v (PrimState m) a
vec
    lst :: Int
lst = Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1
    loop :: Int -> m ()
loop Int
i | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
lst  = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
           | Bool
otherwise = do Int
j <- (Int, Int) -> g -> m Int
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). StatefulGen g m => (Int, Int) -> g -> m Int
uniformRM (Int
i,Int
lst) g
gen
                            v (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
M.unsafeSwap v (PrimState m) a
vec Int
i Int
j
                            Int -> m ()
loop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)


sqr :: Double -> Double
sqr :: Double -> Double
sqr Double
x = Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x
{-# INLINE sqr #-}

pkgError :: String -> String -> a
pkgError :: forall a. String -> String -> a
pkgError String
func String
msg = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"System.Random.MWC.Distributions." String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
func String -> String -> String
forall a. [a] -> [a] -> [a]
++
                            String
": " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
msg

-- | Random variate generator for Binomial distribution. Will throw
-- exception when parameters are out range.
--
-- The probability of getting exactly k successes in n trials is
-- given by the probability mass function:
--
-- \[
-- f(k;n,p) = \Pr(X = k) = \binom n k  p^k(1-p)^{n-k}
-- \]
binomial :: forall g m . StatefulGen g m
         => Int               -- ^ Number of trials, must be positive.
         -> Double            -- ^ Probability of success \(p \in [0,1]\)
         -> g                 -- ^ Generator
         -> m Int
{-# INLINE binomial #-}
binomial :: forall g (m :: * -> *).
StatefulGen g m =>
Int -> Double -> g -> m Int
binomial Int
n Double
p g
gen
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0             = String -> String -> m Int
forall a. String -> String -> a
pkgError String
"binomial" String
"number of trials must be positive"
  | Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0.0 Bool -> Bool -> Bool
|| Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1.0 = String -> String -> m Int
forall a. String -> String -> a
pkgError String
"binomial" String
"probability must be >= 0 and <= 1"
  | Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0.0 = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0
  | Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1.0 = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
n
  | Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0.5 = if
      | Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
inv_thr -> Int -> Double -> g -> m Int
forall g (m :: * -> *).
StatefulGen g m =>
Int -> Double -> g -> m Int
binomialInv Int
n Double
p g
gen
      | Bool
otherwise                    -> Int -> Double -> g -> m Int
forall g (m :: * -> *).
StatefulGen g m =>
Int -> Double -> g -> m Int
binomialTPE Int
n Double
p g
gen
  | Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0.5  = do
      Int
ix <- case Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p of
        Double
p' | Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p' Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
inv_thr -> Int -> Double -> g -> m Int
forall g (m :: * -> *).
StatefulGen g m =>
Int -> Double -> g -> m Int
binomialInv Int
n Double
p' g
gen
           | Bool
otherwise                     -> Int -> Double -> g -> m Int
forall g (m :: * -> *).
StatefulGen g m =>
Int -> Double -> g -> m Int
binomialTPE Int
n Double
p' g
gen
      Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$! Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ix
    -- Reachable when p is NaN
  | Bool
otherwise = String -> String -> m Int
forall a. String -> String -> a
pkgError String
"binomial" String
"probability must be >= 0 and <= 1"
  where
    -- Threshold for preferring the BINV algorithm / inverse cdf
    -- logic. The paper suggests 10, Ranlib uses 30, R uses 30, Rust uses
    -- 10 and GSL uses 14.
    inv_thr :: Double
inv_thr = Double
10

-- Binomial-Triangle-Parallelogram-Exponential algorithm (BTPE)
-- described in Kachitvichyanukul1988
binomialTPE :: forall g m . StatefulGen g m => Int -> Double -> g -> m Int
{-# INLINE binomialTPE #-}
binomialTPE :: forall g (m :: * -> *).
StatefulGen g m =>
Int -> Double -> g -> m Int
binomialTPE Int
n Double
p g
g = m Int
loop
  where
    -- Main accept/reject loop
    loop :: m Int
loop = do
      Double
u <- (Double, Double) -> g -> m Double
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *).
StatefulGen g m =>
(Double, Double) -> g -> m Double
uniformRM (Double
0.0, Double
p4) g
g
      Double
v <- g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
g
      Double -> Double -> m Int
selectArea Double
u Double
v
    -- Acceptance / rejection of sample [step 5]
    acceptReject :: Int -> Double -> m Int
    acceptReject :: Int -> Double -> m Int
acceptReject !Int
ix !Double
v
      | Double
var Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
accept = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
ix
      | Bool
otherwise     = m Int
loop
      where
        var :: Double
var    = Double -> Double
forall a. Floating a => a -> a
log Double
v
        accept :: Double
accept = Int -> Double
forall a. Integral a => a -> Double
logFactorial Int
bigM Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a. Integral a => a -> Double
logFactorial (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bigM) Double -> Double -> Double
forall a. Num a => a -> a -> a
-
                 Int -> Double
forall a. Integral a => a -> Double
logFactorial Int
ix Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
forall a. Integral a => a -> Double
logFactorial (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ix) Double -> Double -> Double
forall a. Num a => a -> a -> a
+
                 Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bigM) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log (Double
p Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
q)
    -- Select area to be used [Steps 1-4]
    selectArea :: Double -> Double -> m Int
    selectArea :: Double -> Double -> m Int
selectArea !Double
u !Double
v
        -- Triangular region
      | Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
p1 = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$! Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Double
xm Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
u
        -- Parallelogram region
      | Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
p2 = do let x :: Double
x = Double
xl Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
c
                         w :: Double
w = Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1.0 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Num a => a -> a
abs (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
xm) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
p1
                     if Double
w Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1 Bool -> Bool -> Bool
|| Double
w Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0
                       then m Int
loop
                       else do let ix :: Int
ix = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x
                               Int -> Double -> m Int
acceptReject Int
ix Double
w
        -- Left tail
      | Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
p3 = case Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Double
xl Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log Double
v Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
lambdaL of
          Int
ix | Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0    -> m Int
loop
             | Bool
otherwise -> do let w :: Double
w = Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p2) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
lambdaL
                               Int -> Double -> m Int
acceptReject Int
ix Double
w
        -- Right tail
      | Bool
otherwise = case Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Double
xr Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
log Double
v Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
lambdaR of
          Int
ix | Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n    -> m Int
loop
             | Bool
otherwise -> do let w :: Double
w = Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p3) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
lambdaR
                               Int -> Double -> m Int
acceptReject Int
ix Double
w
    ----------------------------------------
    -- Constants used in algorithm. See [Step 0]
    q :: Double
q    = Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p
    np :: Double
np   = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p
    ffm :: Double
ffm  = Double
np Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
p
    bigM :: Int
bigM = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
ffm
    -- Half integer mean (tip of triangle)
    xm :: Double
xm   = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bigM Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5
    -- p1: the distance to the left and right edges of the triangle
    -- region below the target distribution; since height=1, also:
    -- area of region (half base * height)
    !p1 :: Double
p1 = let npq :: Double
npq  = Double
np Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
q
          in Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double
2.195 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt Double
npq Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
4.6 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
q) :: Int) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5
    xl :: Double
xl = Double
xm Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p1   -- Left edge of triangle
    xr :: Double
xr = Double
xm Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
p1   -- Right edge of triangle
    c :: Double
c  = Double
0.134 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
20.5 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
15.3 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bigM)
    -- p1 + area of parallelogram region
    !p2 :: Double
p2 = Double
p1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1.0 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
c)
    --
    lambdaL :: Double
lambdaL = let al :: Double
al = (Double
ffm Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
xl) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
ffm Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
xl Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p)
              in Double
al Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1.0 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
al)
    lambdaR :: Double
lambdaR = let ar :: Double
ar = (Double
xr Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
ffm) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
xr Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
q)
              in Double
ar Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1.0 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
ar)
    -- p2 + area of left tail
    !p3 :: Double
p3 = Double
p2 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
c Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
lambdaL
    -- p3 + area of right tail
    !p4 :: Double
p4 = Double
p3 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
c Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
lambdaR


-- Compute binomial variate using inversion method (BINV in
-- Kachitvichyanukul1988)
binomialInv :: StatefulGen g m => Int -> Double -> g -> m Int
{-# INLINE binomialInv #-}
binomialInv :: forall g (m :: * -> *).
StatefulGen g m =>
Int -> Double -> g -> m Int
binomialInv Int
n Double
p g
g = do
  Double
u <- g -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
g
  Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$! Int -> Double -> Double -> Int
invertBinomial Int
n Double
p Double
u

-- This function is defined on top level to avoid inlining it since it's rather
-- large and we don't need specializations since it's monomorphic anyway
invertBinomial
  :: Int    -- N of trials
  -> Double -- probability of success
  -> Double -- Output of PRNG
  -> Int
invertBinomial :: Int -> Double -> Double -> Int
invertBinomial !Int
n !Double
p !Double
u0 = Double -> Double -> Int -> Int
invert (Double
qDouble -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^Int
n) Double
u0 Int
0
  where
    -- We forcing s&a in order to avoid allocating thunks. Those are
    -- more expensive than computing them unconditionally
    q :: Double
q  = Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p
    !s :: Double
s = Double
p Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
q
    !a :: Double
a = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
s
    --
    invert :: Double -> Double -> Int -> Int
invert !Double
r !Double
u !Int
x
      | Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
r    = Int
x
      | Bool
otherwise = Double -> Double -> Int -> Int
invert Double
r' Double
u' Int
x'
      where
        u' :: Double
u' = Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
r
        x' :: Int
x' = Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        r' :: Double
r' = Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
* ((Double
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x') Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
s)


-- $references
--
-- * Doornik, J.A. (2005) An improved ziggurat method to generate
--   normal random samples. Mimeo, Nuffield College, University of
--   Oxford.  <http://www.doornik.com/research/ziggurat.pdf>
--
-- * Thomas, D.B.; Leong, P.G.W.; Luk, W.; Villasenor, J.D.
--   (2007). Gaussian random number generators.
--   /ACM Computing Surveys/ 39(4).
--   <http://www.cse.cuhk.edu.hk/~phwl/mt/public/archives/papers/grng_acmcs07.pdf>
--
-- * Kachitvichyanukul, V. and Schmeiser, B. W.  Binomial Random
--   Variate Generation.  Communications of the ACM, 31, 2 (February,
--   1988) 216. <https://dl.acm.org/doi/pdf/10.1145/42372.42381>
--   Here's an example of how the algorithm's sampling regions look
--   ![Something](docs/RecreateFigure.svg)