-- |
-- Module     : Simulation.Aivika.Trans.Generator.Primitive
-- Copyright  : Copyright (c) 2009-2017, David Sorokin <david.sorokin@gmail.com>
-- License    : BSD3
-- Maintainer : David Sorokin <david.sorokin@gmail.com>
-- Stability  : experimental
-- Tested with: GHC 8.0.1
--
-- This helper module defines primitives for generating random numbers.
--
module Simulation.Aivika.Trans.Generator.Primitive where

import Control.Monad
import Control.Monad.Trans

import Simulation.Aivika.Generator (DiscretePDF)

-- | Generate an uniform random number with the specified minimum and maximum.
generateUniform01 :: Monad m
                     => m Double
                     -- ^ the uniform random number ~ U (0, 1)
                     -> Double
                     -- ^ minimum
                     -> Double
                     -- ^ maximum
                     -> m Double
{-# INLINE generateUniform01 #-}
generateUniform01 :: m Double -> Double -> Double -> m Double
generateUniform01 m Double
g Double
min Double
max =
  do Double
x <- m Double
g
     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
min Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
max Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
min)

-- | Generate an uniform random number with the specified minimum and maximum.
generateUniformInt01 :: Monad m
                        => m Double
                        -- ^ the uniform random number ~ U (0, 1)
                        -> Int
                        -- ^ minimum
                        -> Int
                        -- ^ maximum
                        -> m Int
{-# INLINE generateUniformInt01 #-}
generateUniformInt01 :: m Double -> Int -> Int -> m Int
generateUniformInt01 m Double
g Int
min Int
max =
  do Double
x <- m Double
g
     let min' :: Double
min' = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
min Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
0.5
         max' :: Double
max' = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
max Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5
         z :: Int
z    = Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
round (Double
min' Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
max' Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
min'))
         z' :: Int
z'   = if Int
z Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
min
                then Int
min
                else if Int
z Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
max
                     then Int
max
                     else Int
z
     Int -> m Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
z'

-- | Generate the triangular random number by the specified minimum, median and maximum.
generateTriangular01 :: Monad m
                        => m Double
                        -- ^ the uniform random number ~ U (0, 1)
                        -> Double
                        -- ^ minimum
                        -> Double
                        -- ^ median
                        -> Double
                        -- ^ maximum
                        -> m Double
{-# INLINE generateTriangular01 #-}
generateTriangular01 :: m Double -> Double -> Double -> Double -> m Double
generateTriangular01 m Double
g Double
min Double
median Double
max =
  do Double
x <- m Double
g
     if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= (Double
median Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
min) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
max Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
min)
       then 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
min Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
sqrt ((Double
median Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
min) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
max Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
min) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x)
       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
$ Double
max Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
sqrt ((Double
max Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
median) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
max Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
min) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
x))

-- | Generate a normal random number by the specified generator, mean and variance.
generateNormal01 :: Monad m
                    => m Double
                    -- ^ the normal random number ~ N (0, 1)
                    -> Double
                    -- ^ mean
                    -> Double
                    -- ^ variance
                    -> m Double
{-# INLINE generateNormal01 #-}
generateNormal01 :: m Double -> Double -> Double -> m Double
generateNormal01 m Double
g Double
mu Double
nu =
  do Double
x <- m Double
g
     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
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
nu Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x

-- | Generate the lognormal random number derived from a normal distribution with
-- the specified generator, mean and variance.
generateLogNormal01 :: Monad m
                       => m Double
                       -- ^ the normal random number ~ N (0, 1)
                       -> Double
                       -- ^ mean
                       -> Double
                       -- ^ variance
                       -> m Double
{-# INLINE generateLogNormal01 #-}
generateLogNormal01 :: m Double -> Double -> Double -> m Double
generateLogNormal01 m Double
g Double
mu Double
nu =
  do Double
x <- m Double
g
     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
exp (Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
nu Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x)

-- | Return the exponential random number with the specified mean.
generateExponential01 :: Monad m
                         => m Double
                         -- ^ the uniform random number ~ U (0, 1)
                         -> Double
                         -- ^ the mean
                         -> m Double
{-# INLINE generateExponential01 #-}
generateExponential01 :: m Double -> Double -> m Double
generateExponential01 m Double
g Double
mu =
  do Double
x <- m Double
g
     Double -> m Double
forall (m :: * -> *) a. Monad m => a -> m a
return (- Double -> Double
forall a. Floating a => a -> a
log Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
mu)

-- | Return the Erlang random number.
generateErlang01 :: Monad m
                    => m Double
                    -- ^ the uniform random number ~ U (0, 1)
                    -> Double
                    -- ^ the scale
                    -> Int
                    -- ^ the shape
                    -> m Double
{-# INLINABLE generateErlang01 #-}
generateErlang01 :: m Double -> Double -> Int -> m Double
generateErlang01 m Double
g Double
beta Int
m =
  do Double
x <- Int -> Double -> m Double
forall t. (Ord t, Num t) => t -> Double -> m Double
loop Int
m Double
1
     Double -> m Double
forall (m :: * -> *) a. Monad m => a -> m a
return (- Double -> Double
forall a. Floating a => a -> a
log Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
beta)
       where loop :: t -> Double -> m Double
loop t
m Double
acc
               | t
m t -> t -> Bool
forall a. Ord a => a -> a -> Bool
< t
0     = [Char] -> m Double
forall a. HasCallStack => [Char] -> a
error [Char]
"Negative shape: generateErlang."
               | t
m t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
0    = Double -> m Double
forall (m :: * -> *) a. Monad m => a -> m a
return Double
acc
               | Bool
otherwise = do Double
x <- m Double
g
                                t -> Double -> m Double
loop (t
m t -> t -> t
forall a. Num a => a -> a -> a
- t
1) (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
acc)

-- | Generate the Poisson random number with the specified mean.
generatePoisson01 :: Monad m
                     => m Double
                     -- ^ the uniform random number ~ U (0, 1)
                     -> Double
                     -- ^ the mean
                     -> m Int
{-# INLINABLE generatePoisson01 #-}
generatePoisson01 :: m Double -> Double -> m Int
generatePoisson01 m Double
g Double
mu =
  do Double
prob0 <- m Double
g
     let loop :: Double -> Double -> t -> m t
loop Double
prob Double
prod t
acc
           | Double
prob Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
prod = t -> m t
forall (m :: * -> *) a. Monad m => a -> m a
return t
acc
           | Bool
otherwise    = Double -> Double -> t -> m t
loop
                            (Double
prob Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
prod)
                            (Double
prod Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
mu Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ t -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (t
acc t -> t -> t
forall a. Num a => a -> a -> a
+ t
1))
                            (t
acc t -> t -> t
forall a. Num a => a -> a -> a
+ t
1)
     Double -> Double -> Int -> m Int
forall (m :: * -> *) t.
(Monad m, Integral t) =>
Double -> Double -> t -> m t
loop Double
prob0 (Double -> Double
forall a. Floating a => a -> a
exp (- Double
mu)) Int
0

-- | Generate a binomial random number with the specified probability and number of trials. 
generateBinomial01 :: Monad m
                      => m Double
                      -- ^ the uniform random number ~ U (0, 1)
                      -> Double 
                      -- ^ the probability
                      -> Int
                      -- ^ the number of trials
                      -> m Int
{-# INLINABLE generateBinomial01 #-}
generateBinomial01 :: m Double -> Double -> Int -> m Int
generateBinomial01 m Double
g Double
prob Int
trials = Int -> Int -> m Int
forall a t. (Ord a, Num a, Num t) => a -> t -> m t
loop Int
trials Int
0 where
  loop :: a -> t -> m t
loop a
n t
acc
    | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0     = [Char] -> m t
forall a. HasCallStack => [Char] -> a
error [Char]
"Negative number of trials: generateBinomial."
    | a
n a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0    = t -> m t
forall (m :: * -> *) a. Monad m => a -> m a
return t
acc
    | Bool
otherwise = do Double
x <- m Double
g
                     if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
prob
                       then a -> t -> m t
loop (a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
1) (t
acc t -> t -> t
forall a. Num a => a -> a -> a
+ t
1)
                       else a -> t -> m t
loop (a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
1) t
acc

-- | Generate a random number from the Gamma distribution using Marsaglia and Tsang method.
generateGamma01 :: Monad m
                   => m Double
                   -- ^ the normal random number ~ N (0,1)
                   -> m Double
                   -- ^ the uniform random number ~ U (0, 1)
                   -> Double
                   -- ^ the shape parameter (kappa) 
                   -> Double
                   -- ^ the scale parameter (theta)
                   -> m Double
{-# INLINABLE generateGamma01 #-}
generateGamma01 :: m Double -> m Double -> Double -> Double -> m Double
generateGamma01 m Double
gn m Double
gu Double
kappa Double
theta
  | Double
kappa Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 = [Char] -> m Double
forall a. HasCallStack => [Char] -> a
error [Char]
"The shape parameter (kappa) must be positive: generateGamma01"
  | Double
kappa Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1  =
    let d :: Double
d = Double
kappa Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
3
        c :: Double
c = 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
d)
        loop :: m Double
loop =
          do Double
z <- m Double
gn
             if Double
z Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= - (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
c)
               then m Double
loop
               else do let v :: Double
v = (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
z) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
3
                       Double
u <- m Double
gu
                       if 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
z Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
z Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
d Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
d Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
d Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
v
                         then m Double
loop
                         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
$ Double
d Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
theta
    in m Double
loop
  | Bool
otherwise  =
    do Double
x <- m Double -> m Double -> Double -> Double -> m Double
forall (m :: * -> *).
Monad m =>
m Double -> m Double -> Double -> Double -> m Double
generateGamma01 m Double
gn m Double
gu (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
kappa) Double
theta
       Double
u <- m Double
gu
       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. Num a => a -> a -> a
* Double
u Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
kappa)

-- | Generate a random number from the Beta distribution.
generateBeta01 :: Monad m
                  => m Double
                  -- ^ the normal random number ~ N (0, 1)
                  -> m Double
                  -- ^ the uniform random number ~ U (0, 1)
                  -> Double
                  -- ^ the shape parameter alpha
                  -> Double
                  -- ^ the shape parameter beta
                  -> m Double
{-# INLINABLE generateBeta01 #-}
generateBeta01 :: m Double -> m Double -> Double -> Double -> m Double
generateBeta01 m Double
gn m Double
gu Double
alpha Double
beta =
  do Double
g1 <- m Double -> m Double -> Double -> Double -> m Double
forall (m :: * -> *).
Monad m =>
m Double -> m Double -> Double -> Double -> m Double
generateGamma01 m Double
gn m Double
gu Double
alpha Double
1
     Double
g2 <- m Double -> m Double -> Double -> Double -> m Double
forall (m :: * -> *).
Monad m =>
m Double -> m Double -> Double -> Double -> m Double
generateGamma01 m Double
gn m Double
gu Double
beta 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
g1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
g1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
g2)

-- | Generate a random number from the Weibull distribution.
generateWeibull01 :: Monad m
                     => m Double
                     -- ^ the uniform random number ~ U (0, 1)
                     -> Double
                     -- ^ shape
                     -> Double
                     -- ^ scale
                     -> m Double
{-# INLINE generateWeibull01 #-}
generateWeibull01 :: m Double -> Double -> Double -> m Double
generateWeibull01 m Double
g Double
alpha Double
beta =
  do Double
x <- m Double
g
     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
beta Double -> Double -> Double
forall a. Num a => a -> a -> a
* (- Double -> Double
forall a. Floating a => a -> a
log Double
x) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
alpha)

-- | Generate a random value from the specified discrete distribution.
generateDiscrete01 :: Monad m
                      => m Double
                      -- ^ the uniform random number ~ U (0, 1)
                      -> DiscretePDF a
                      -- ^ a discrete probability density function
                      -> m a
{-# INLINABLE generateDiscrete01 #-}
generateDiscrete01 :: m Double -> DiscretePDF a -> m a
generateDiscrete01 m Double
g []   = [Char] -> m a
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty PDF: generateDiscrete01"
generateDiscrete01 m Double
g DiscretePDF a
dpdf =
  do Double
x <- m Double
g
     let loop :: Double -> [(p, Double)] -> p
loop Double
acc [(p
a, Double
p)] = p
a
         loop Double
acc ((p
a, Double
p) : [(p, Double)]
dpdf) =
           if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
acc Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
p
           then p
a
           else Double -> [(p, Double)] -> p
loop (Double
acc Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
p) [(p, Double)]
dpdf
     a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> m a) -> a -> m a
forall a b. (a -> b) -> a -> b
$ Double -> DiscretePDF a -> a
forall p. Double -> [(p, Double)] -> p
loop Double
0 DiscretePDF a
dpdf