{-# 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
      -- ** 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

-- 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 :: 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 (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 :: 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
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 (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 (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 (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 :: 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 (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 :: 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 (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 :: 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 (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 (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 (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 :: 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 (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 :: 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 (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 (m :: * -> *) a. Monad m => a -> m a
return (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$! Double -> Int
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 :: 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 (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 :: 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 (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 :: 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)
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 (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 (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 (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 :: 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 :: 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 (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 :: 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 :: 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 :: 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 (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 :: 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 (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 (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
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 :: 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

-- $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>