{-# 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 :: forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
normal Double
m Double
s g
gen = do
  Double
x <- forall g (m :: * -> *). StatefulGen g m => g -> m Double
standard g
gen
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Double
m forall a. Num a => a -> a -> a
+ Double
s 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  <- (forall a. Num a => a -> a -> a
subtract Double
1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Num a => a -> a -> a
*Double
2)) forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
      Word32
ri <- forall a g (m :: * -> *). (Uniform a, StatefulGen g m) => g -> m a
uniformM g
gen
      let i :: Int
i  = forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word32
ri :: Word32) forall a. Bits a => a -> a -> a
.&. Word32
127)
          bi :: Double
bi = forall a. Unbox a => Vector a -> Int -> a
I.unsafeIndex Vector Double
blocks Int
i
          bj :: Double
bj = forall a. Unbox a => Vector a -> Int -> a
I.unsafeIndex Vector Double
blocks (Int
iforall a. Num a => a -> a -> a
+Int
1)
      case () of
        ()
_| forall a. Num a => a -> a
abs Double
u forall a. Ord a => a -> a -> Bool
< forall a. Unbox a => Vector a -> Int -> a
I.unsafeIndex Vector Double
ratios Int
i -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Double
u forall a. Num a => a -> a -> a
* Double
bi
         | Int
i forall a. Eq a => a -> a -> Bool
== Int
0                         -> Bool -> m Double
normalTail (Double
u forall a. Ord a => a -> a -> Bool
< Double
0)
         | Bool
otherwise                      -> do
             let x :: Double
x  = Double
u forall a. Num a => a -> a -> a
* Double
bi
                 xx :: Double
xx = Double
x forall a. Num a => a -> a -> a
* Double
x
                 d :: Double
d  = forall a. Floating a => a -> a
exp (-Double
0.5 forall a. Num a => a -> a -> a
* (Double
bi forall a. Num a => a -> a -> a
* Double
bi forall a. Num a => a -> a -> a
- Double
xx))
                 e :: Double
e  = forall a. Floating a => a -> a
exp (-Double
0.5 forall a. Num a => a -> a -> a
* (Double
bj forall a. Num a => a -> a -> a
* Double
bj forall a. Num a => a -> a -> a
- Double
xx))
             Double
c <- forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
             if Double
e forall a. Num a => a -> a -> a
+ Double
c forall a. Num a => a -> a -> a
* (Double
d forall a. Num a => a -> a -> a
- Double
e) forall a. Ord a => a -> a -> Bool
< Double
1
               then 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 <- ((forall a. Fractional a => a -> a -> a
/Double
rNorm) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
log) forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
              Double
y <- forall a. Floating a => a -> a
log              forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
              if Double
y forall a. Num a => a -> a -> a
* (-Double
2) forall a. Ord a => a -> a -> Bool
< Double
x forall a. Num a => a -> a -> a
* Double
x
                then m Double
tailing
                else forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! if Bool
neg then Double
x forall a. Num a => a -> a -> a
- Double
rNorm else Double
rNorm 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 = (forall a. Unbox a => Vector a -> a -> Vector a
`I.snoc` Double
0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Unbox a => a -> Vector a -> Vector a
I.cons (Double
vforall a. Fractional a => a -> a -> a
/Double
f) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Unbox a => a -> Vector a -> Vector a
I.cons Double
rNorm forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Unbox a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
I.unfoldrN Int
126 T -> Maybe (Double, T)
go 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 (forall a. Floating a => a -> a
exp (-Double
0.5 forall a. Num a => a -> a -> a
* Double
h forall a. Num a => a -> a -> a
* Double
h))
                     h :: Double
h  = forall a. Floating a => a -> a
sqrt (-Double
2 forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
log (Double
v forall a. Fractional a => a -> a -> a
/ Double
b forall a. Num a => a -> a -> a
+ Double
g))
                 in forall a. a -> Maybe a
Just (Double
h, T
u)
    v :: Double
v = Double
9.91256303526217e-3
    f :: Double
f = forall a. Floating a => a -> a
exp (-Double
0.5 forall a. Num a => a -> a -> a
* Double
rNorm 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 = forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
I.zipWith forall a. Fractional a => a -> a -> a
(/) (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 <- forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! - forall a. Floating a => a -> a
log Double
x 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 forall a. Num a => a -> a -> a
- Double
a
  Double
p <- forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Double
a forall a. Num a => a -> a -> a
- forall a. Floating a => a -> a
log ( (Double
1 forall a. Num a => a -> a -> a
- Double
p) forall a. Num a => a -> a -> a
+ Double
pforall a. Num a => a -> a -> a
*forall a. Floating a => a -> a
exp(-Double
scaleforall a. Num a => a -> a -> a
*Double
delta)) 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 forall a. Ord a => a -> a -> Bool
<= Double
0    = 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     <- forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
        let cont :: Bool
cont =  Double
u forall a. Ord a => a -> a -> Bool
> Double
1 forall a. Num a => a -> a -> a
- Double
0.331 forall a. Num a => a -> a -> a
* Double -> Double
sqr (Double -> Double
sqr Double
x)
                 Bool -> Bool -> Bool
&& forall a. Floating a => a -> a
log Double
u forall a. Ord a => a -> a -> Bool
> Double
0.5 forall a. Num a => a -> a -> a
* Double -> Double
sqr Double
x forall a. Num a => a -> a -> a
+ Double
a1 forall a. Num a => a -> a -> a
* (Double
1 forall a. Num a => a -> a -> a
- Double
v forall a. Num a => a -> a -> a
+ forall a. Floating a => a -> a
log Double
v) -- Rarely evaluated
        case () of
          ()
_| Bool
cont      -> m Double
mainloop
           | Double
a forall a. Ord a => a -> a -> Bool
>= Double
1    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Double
a1 forall a. Num a => a -> a -> a
* Double
v forall a. Num a => a -> a -> a
* Double
b
           | Bool
otherwise -> do Double
y <- forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
                             forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Double
y forall a. Floating a => a -> a -> a
** (Double
1 forall a. Fractional a => a -> a -> a
/ Double
a) forall a. Num a => a -> a -> a
* Double
a1 forall a. Num a => a -> a -> a
* Double
v forall a. Num a => a -> a -> a
* Double
b
      -- inner loop
      innerloop :: m T
innerloop = do
        Double
x <- forall g (m :: * -> *). StatefulGen g m => g -> m Double
standard g
gen
        case Double
1 forall a. Num a => a -> a -> a
+ Double
a2forall a. Num a => a -> a -> a
*Double
x of
          Double
v | Double
v forall a. Ord a => a -> a -> Bool
<= Double
0    -> m T
innerloop
            | Bool
otherwise -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Double -> Double -> T
T Double
x (Double
vforall a. Num a => a -> a -> a
*Double
vforall a. Num a => a -> a -> a
*Double
v)
      -- constants
      a' :: Double
a' = if Double
a forall a. Ord a => a -> a -> Bool
< Double
1 then Double
a forall a. Num a => a -> a -> a
+ Double
1 else Double
a
      a1 :: Double
a1 = Double
a' forall a. Num a => a -> a -> a
- Double
1forall a. Fractional a => a -> a -> a
/Double
3
      a2 :: Double
a2 = Double
1 forall a. Fractional a => a -> a -> a
/ forall a. Floating a => a -> a
sqrt(Double
9 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 forall a. Ord a => a -> a -> Bool
<= Int
0    = forall a. String -> String -> a
pkgError String
"chiSquare" String
"number of degrees of freedom must be positive"
  | Bool
otherwise = do Double
x <- forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
gamma (Double
0.5 forall a. Num a => a -> a -> a
* forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) Double
1 g
gen
                   forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Double
2 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 forall a. Eq a => a -> a -> Bool
== Double
1          = forall (m :: * -> *) a. Monad m => a -> m a
return Int
0
  | Double
p forall a. Ord a => a -> a -> Bool
>  Double
0 Bool -> Bool -> Bool
&& Double
p forall a. Ord a => a -> a -> Bool
< Double
1 = do Double
q <- 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.
                         forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a b. (RealFrac a, Integral b) => a -> b
floor forall a b. (a -> b) -> a -> b
$ forall a. Floating a => a -> a
log Double
q forall a. Fractional a => a -> a -> a
/ forall a. Floating a => a -> a
log (Double
1 forall a. Num a => a -> a -> a
- Double
p)
  | Bool
otherwise       = 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 <- forall g (m :: * -> *). StatefulGen g m => Double -> g -> m Int
geometric0 Double
p g
gen
                      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Int
n 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 <- forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
gamma Double
a Double
1 g
gen
  Double
y <- forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
gamma Double
b Double
1 g
gen
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Double
x forall a. Fractional a => a -> a -> a
/ (Double
xforall 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' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\Double
x -> 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 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Num a => a -> a -> a
(+) Double
0 t Double
t'
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (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 = (forall a. Ord a => a -> a -> Bool
<Double
p) forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` 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
    | forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v Double
v = forall a. String -> String -> a
pkgError String
"categorical" String
"empty weights!"
    | Bool
otherwise = do
        let cv :: v Double
cv  = forall (v :: * -> *) a. Vector v a => (a -> a -> a) -> v a -> v a
G.scanl1' forall a. Num a => a -> a -> a
(+) v Double
v
        Double
p <- (forall (v :: * -> *) a. Vector v a => v a -> a
G.last v Double
cv forall a. Num a => a -> a -> a
*) forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDoublePositive01M g
gen
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! case forall (v :: * -> *) a.
Vector v a =>
(a -> Bool) -> v a -> Maybe Int
G.findIndex (forall a. Ord a => a -> a -> Bool
>=Double
p) v Double
cv of
                    Just Int
i  -> Int
i
                    Maybe Int
Nothing -> 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
  | forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v Double
v  = forall a. String -> String -> a
pkgError String
"logCategorical" String
"empty weights!"
  | Bool
otherwise = forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical (forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map (forall a. Floating a => a -> a
exp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Num a => a -> a -> a
subtract Double
m) v Double
v) g
gen
  where
    m :: Double
m = 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 forall a. Ord a => a -> a -> Bool
< Int
0     = forall a. String -> String -> a
pkgError String
"uniformPermutation" String
"size must be >=0"
  | Bool
otherwise = forall g (m :: * -> *) (v :: * -> *) a.
(StatefulGen g m, PrimMonad m, Vector v a) =>
v a -> g -> m (v a)
uniformShuffle (forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
G.generate Int
n 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
  | forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
vec forall a. Ord a => a -> a -> Bool
<= Int
1 = forall (m :: * -> *) a. Monad m => a -> m a
return v a
vec
  | Bool
otherwise         = do
      Mutable v (PrimState m) a
mvec <- forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
G.thaw v a
vec
      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
      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
  | forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
M.length v (PrimState m) a
vec forall a. Ord a => a -> a -> Bool
<= Int
1 = forall (m :: * -> *) a. Monad m => a -> m a
return ()
  | Bool
otherwise         = Int -> m ()
loop Int
0
  where
    n :: Int
n   = forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
M.length v (PrimState m) a
vec
    lst :: Int
lst = Int
nforall a. Num a => a -> a -> a
-Int
1
    loop :: Int -> m ()
loop Int
i | Int
i forall a. Eq a => a -> a -> Bool
== Int
lst  = forall (m :: * -> *) a. Monad m => a -> m a
return ()
           | Bool
otherwise = do Int
j <- forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
uniformRM (Int
i,Int
lst) g
gen
                            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
iforall a. Num a => a -> a -> a
+Int
1)


sqr :: Double -> Double
sqr :: Double -> Double
sqr Double
x = Double
x 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 = forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"System.Random.MWC.Distributions." forall a. [a] -> [a] -> [a]
++ String
func forall a. [a] -> [a] -> [a]
++
                            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>