{-# LANGUAGE RankNTypes #-}
module Simulation.Aivika.Generator
(Generator(..),
GeneratorType(..),
DiscretePDF(..),
newGenerator,
newRandomGenerator,
newRandomGenerator01) where
import System.Random
import qualified System.Random.MWC as MWC
import Data.IORef
import Data.Word
import Data.Vector
import Data.Functor
type DiscretePDF a = [(a, Double)]
data Generator =
Generator { Generator -> Double -> Double -> IO Double
generateUniform :: Double -> Double -> IO Double,
Generator -> Int -> Int -> IO Int
generateUniformInt :: Int -> Int -> IO Int,
Generator -> Double -> Double -> Double -> IO Double
generateTriangular :: Double -> Double -> Double -> IO Double,
Generator -> Double -> Double -> IO Double
generateNormal :: Double -> Double -> IO Double,
Generator -> Double -> Double -> IO Double
generateLogNormal :: Double -> Double -> IO Double,
Generator -> Double -> IO Double
generateExponential :: Double -> IO Double,
Generator -> Double -> Int -> IO Double
generateErlang :: Double -> Int -> IO Double,
Generator -> Double -> IO Int
generatePoisson :: Double -> IO Int,
Generator -> Double -> Int -> IO Int
generateBinomial :: Double -> Int -> IO Int,
Generator -> Double -> Double -> IO Double
generateGamma :: Double -> Double -> IO Double,
Generator -> Double -> Double -> IO Double
generateBeta :: Double -> Double -> IO Double,
Generator -> Double -> Double -> IO Double
generateWeibull :: Double -> Double -> IO Double,
Generator -> forall a. DiscretePDF a -> IO a
generateDiscrete :: forall a. DiscretePDF a -> IO a,
Generator -> IO Int
generateSequenceNo :: IO Int
}
generateUniform01 :: IO Double
-> Double
-> Double
-> IO Double
generateUniform01 :: IO Double -> Double -> Double -> IO Double
generateUniform01 IO Double
g Double
min Double
max =
do Double
x <- IO Double
g
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Double
min forall a. Num a => a -> a -> a
+ Double
x forall a. Num a => a -> a -> a
* (Double
max forall a. Num a => a -> a -> a
- Double
min)
generateUniformInt01 :: IO Double
-> Int
-> Int
-> IO Int
generateUniformInt01 :: IO Double -> Int -> Int -> IO Int
generateUniformInt01 IO Double
g Int
min Int
max =
do Double
x <- IO Double
g
let min' :: Double
min' = forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
min forall a. Num a => a -> a -> a
- Double
0.5
max' :: Double
max' = forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
max forall a. Num a => a -> a -> a
+ Double
0.5
z :: Int
z = forall a b. (RealFrac a, Integral b) => a -> b
round (Double
min' forall a. Num a => a -> a -> a
+ Double
x forall a. Num a => a -> a -> a
* (Double
max' forall a. Num a => a -> a -> a
- Double
min'))
z' :: Int
z' = if Int
z forall a. Ord a => a -> a -> Bool
< Int
min
then Int
min
else if Int
z forall a. Ord a => a -> a -> Bool
> Int
max
then Int
max
else Int
z
forall (m :: * -> *) a. Monad m => a -> m a
return Int
z'
generateTriangular01 :: IO Double
-> Double
-> Double
-> Double
-> IO Double
generateTriangular01 :: IO Double -> Double -> Double -> Double -> IO Double
generateTriangular01 IO Double
g Double
min Double
median Double
max =
do Double
x <- IO Double
g
if Double
x forall a. Ord a => a -> a -> Bool
<= (Double
median forall a. Num a => a -> a -> a
- Double
min) forall a. Fractional a => a -> a -> a
/ (Double
max forall a. Num a => a -> a -> a
- Double
min)
then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Double
min forall a. Num a => a -> a -> a
+ forall a. Floating a => a -> a
sqrt ((Double
median forall a. Num a => a -> a -> a
- Double
min) forall a. Num a => a -> a -> a
* (Double
max forall a. Num a => a -> a -> a
- Double
min) forall a. Num a => a -> a -> a
* Double
x)
else forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Double
max forall a. Num a => a -> a -> a
- forall a. Floating a => a -> a
sqrt ((Double
max forall a. Num a => a -> a -> a
- Double
median) forall a. Num a => a -> a -> a
* (Double
max forall a. Num a => a -> a -> a
- Double
min) forall a. Num a => a -> a -> a
* (Double
1 forall a. Num a => a -> a -> a
- Double
x))
newNormalGenerator01 :: IO Double
-> IO (IO Double)
newNormalGenerator01 :: IO Double -> IO (IO Double)
newNormalGenerator01 IO Double
g =
do IORef Double
nextRef <- forall a. a -> IO (IORef a)
newIORef Double
0.0
IORef Bool
flagRef <- forall a. a -> IO (IORef a)
newIORef Bool
False
IORef Double
xi1Ref <- forall a. a -> IO (IORef a)
newIORef Double
0.0
IORef Double
xi2Ref <- forall a. a -> IO (IORef a)
newIORef Double
0.0
IORef Double
psiRef <- forall a. a -> IO (IORef a)
newIORef Double
0.0
let loop :: IO ()
loop =
do Double
psi <- forall a. IORef a -> IO a
readIORef IORef Double
psiRef
if (Double
psi forall a. Ord a => a -> a -> Bool
>= Double
1.0) Bool -> Bool -> Bool
|| (Double
psi forall a. Eq a => a -> a -> Bool
== Double
0.0)
then do Double
g1 <- IO Double
g
Double
g2 <- IO Double
g
let xi1 :: Double
xi1 = Double
2.0 forall a. Num a => a -> a -> a
* Double
g1 forall a. Num a => a -> a -> a
- Double
1.0
xi2 :: Double
xi2 = Double
2.0 forall a. Num a => a -> a -> a
* Double
g2 forall a. Num a => a -> a -> a
- Double
1.0
psi :: Double
psi = Double
xi1 forall a. Num a => a -> a -> a
* Double
xi1 forall a. Num a => a -> a -> a
+ Double
xi2 forall a. Num a => a -> a -> a
* Double
xi2
forall a. IORef a -> a -> IO ()
writeIORef IORef Double
xi1Ref Double
xi1
forall a. IORef a -> a -> IO ()
writeIORef IORef Double
xi2Ref Double
xi2
forall a. IORef a -> a -> IO ()
writeIORef IORef Double
psiRef Double
psi
IO ()
loop
else forall a. IORef a -> a -> IO ()
writeIORef IORef Double
psiRef forall a b. (a -> b) -> a -> b
$ forall a. Floating a => a -> a
sqrt (- Double
2.0 forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
log Double
psi forall a. Fractional a => a -> a -> a
/ Double
psi)
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
do Bool
flag <- forall a. IORef a -> IO a
readIORef IORef Bool
flagRef
if Bool
flag
then do forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
flagRef Bool
False
forall a. IORef a -> IO a
readIORef IORef Double
nextRef
else do forall a. IORef a -> a -> IO ()
writeIORef IORef Double
xi1Ref Double
0.0
forall a. IORef a -> a -> IO ()
writeIORef IORef Double
xi2Ref Double
0.0
forall a. IORef a -> a -> IO ()
writeIORef IORef Double
psiRef Double
0.0
IO ()
loop
Double
xi1 <- forall a. IORef a -> IO a
readIORef IORef Double
xi1Ref
Double
xi2 <- forall a. IORef a -> IO a
readIORef IORef Double
xi2Ref
Double
psi <- forall a. IORef a -> IO a
readIORef IORef Double
psiRef
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
flagRef Bool
True
forall a. IORef a -> a -> IO ()
writeIORef IORef Double
nextRef forall a b. (a -> b) -> a -> b
$ Double
xi2 forall a. Num a => a -> a -> a
* Double
psi
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Double
xi1 forall a. Num a => a -> a -> a
* Double
psi
generateExponential01 :: IO Double
-> Double
-> IO Double
generateExponential01 :: IO Double -> Double -> IO Double
generateExponential01 IO Double
g Double
mu =
do Double
x <- IO Double
g
forall (m :: * -> *) a. Monad m => a -> m a
return (- forall a. Floating a => a -> a
log Double
x forall a. Num a => a -> a -> a
* Double
mu)
generateErlang01 :: IO Double
-> Double
-> Int
-> IO Double
generateErlang01 :: IO Double -> Double -> Int -> IO Double
generateErlang01 IO Double
g Double
beta Int
m =
do Double
x <- forall {t}. (Ord t, Num t) => t -> Double -> IO Double
loop Int
m Double
1
forall (m :: * -> *) a. Monad m => a -> m a
return (- forall a. Floating a => a -> a
log Double
x forall a. Num a => a -> a -> a
* Double
beta)
where loop :: t -> Double -> IO Double
loop t
m Double
acc
| t
m forall a. Ord a => a -> a -> Bool
< t
0 = forall a. HasCallStack => [Char] -> a
error [Char]
"Negative shape: generateErlang."
| t
m forall a. Eq a => a -> a -> Bool
== t
0 = forall (m :: * -> *) a. Monad m => a -> m a
return Double
acc
| Bool
otherwise = do Double
x <- IO Double
g
t -> Double -> IO Double
loop (t
m forall a. Num a => a -> a -> a
- t
1) (Double
x forall a. Num a => a -> a -> a
* Double
acc)
generatePoisson01 :: IO Double
-> Double
-> IO Int
generatePoisson01 :: IO Double -> Double -> IO Int
generatePoisson01 IO Double
g Double
mu =
do Double
prob0 <- IO Double
g
let loop :: Double -> Double -> t -> m t
loop Double
prob Double
prod t
acc
| Double
prob forall a. Ord a => a -> a -> Bool
<= Double
prod = forall (m :: * -> *) a. Monad m => a -> m a
return t
acc
| Bool
otherwise = Double -> Double -> t -> m t
loop
(Double
prob forall a. Num a => a -> a -> a
- Double
prod)
(Double
prod forall a. Num a => a -> a -> a
* Double
mu forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral (t
acc forall a. Num a => a -> a -> a
+ t
1))
(t
acc forall a. Num a => a -> a -> a
+ t
1)
forall {m :: * -> *} {t}.
(Monad m, Integral t) =>
Double -> Double -> t -> m t
loop Double
prob0 (forall a. Floating a => a -> a
exp (- Double
mu)) Int
0
generateBinomial01 :: IO Double
-> Double
-> Int
-> IO Int
generateBinomial01 :: IO Double -> Double -> Int -> IO Int
generateBinomial01 IO Double
g Double
prob Int
trials = forall {t} {t}. (Ord t, Num t, Num t) => t -> t -> IO t
loop Int
trials Int
0 where
loop :: t -> t -> IO t
loop t
n t
acc
| t
n forall a. Ord a => a -> a -> Bool
< t
0 = forall a. HasCallStack => [Char] -> a
error [Char]
"Negative number of trials: generateBinomial."
| t
n forall a. Eq a => a -> a -> Bool
== t
0 = forall (m :: * -> *) a. Monad m => a -> m a
return t
acc
| Bool
otherwise = do Double
x <- IO Double
g
if Double
x forall a. Ord a => a -> a -> Bool
<= Double
prob
then t -> t -> IO t
loop (t
n forall a. Num a => a -> a -> a
- t
1) (t
acc forall a. Num a => a -> a -> a
+ t
1)
else t -> t -> IO t
loop (t
n forall a. Num a => a -> a -> a
- t
1) t
acc
generateGamma01 :: IO Double
-> IO Double
-> Double
-> Double
-> IO Double
generateGamma01 :: IO Double -> IO Double -> Double -> Double -> IO Double
generateGamma01 IO Double
gn IO Double
gu Double
kappa Double
theta
| Double
kappa forall a. Ord a => a -> a -> Bool
<= Double
0 = forall a. HasCallStack => [Char] -> a
error [Char]
"The shape parameter (kappa) must be positive: generateGamma01"
| Double
kappa forall a. Ord a => a -> a -> Bool
> Double
1 =
let d :: Double
d = Double
kappa forall a. Num a => a -> a -> a
- Double
1 forall a. Fractional a => a -> a -> a
/ Double
3
c :: Double
c = 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
d)
loop :: IO Double
loop =
do Double
z <- IO Double
gn
if Double
z forall a. Ord a => a -> a -> Bool
<= - (Double
1 forall a. Fractional a => a -> a -> a
/ Double
c)
then IO Double
loop
else do let v :: Double
v = (Double
1 forall a. Num a => a -> a -> a
+ Double
c forall a. Num a => a -> a -> a
* Double
z) forall a. Floating a => a -> a -> a
** Double
3
Double
u <- IO Double
gu
if 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
z forall a. Num a => a -> a -> a
* Double
z forall a. Num a => a -> a -> a
+ Double
d forall a. Num a => a -> a -> a
- Double
d forall a. Num a => a -> a -> a
* Double
v forall a. Num a => a -> a -> a
+ Double
d forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
log Double
v
then IO Double
loop
else forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Double
d forall a. Num a => a -> a -> a
* Double
v forall a. Num a => a -> a -> a
* Double
theta
in IO Double
loop
| Bool
otherwise =
do Double
x <- IO Double -> IO Double -> Double -> Double -> IO Double
generateGamma01 IO Double
gn IO Double
gu (Double
1 forall a. Num a => a -> a -> a
+ Double
kappa) Double
theta
Double
u <- IO Double
gu
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Double
x forall a. Num a => a -> a -> a
* Double
u forall a. Floating a => a -> a -> a
** (Double
1 forall a. Fractional a => a -> a -> a
/ Double
kappa)
generateBeta01 :: IO Double
-> IO Double
-> Double
-> Double
-> IO Double
generateBeta01 :: IO Double -> IO Double -> Double -> Double -> IO Double
generateBeta01 IO Double
gn IO Double
gu Double
alpha Double
beta =
do Double
g1 <- IO Double -> IO Double -> Double -> Double -> IO Double
generateGamma01 IO Double
gn IO Double
gu Double
alpha Double
1
Double
g2 <- IO Double -> IO Double -> Double -> Double -> IO Double
generateGamma01 IO Double
gn IO Double
gu Double
beta Double
1
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Double
g1 forall a. Fractional a => a -> a -> a
/ (Double
g1 forall a. Num a => a -> a -> a
+ Double
g2)
generateWeibull01 :: IO Double
-> Double
-> Double
-> IO Double
generateWeibull01 :: IO Double -> Double -> Double -> IO Double
generateWeibull01 IO Double
g Double
alpha Double
beta =
do Double
x <- IO Double
g
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Double
beta forall a. Num a => a -> a -> a
* (- forall a. Floating a => a -> a
log Double
x) forall a. Floating a => a -> a -> a
** (Double
1 forall a. Fractional a => a -> a -> a
/ Double
alpha)
generateDiscrete01 :: IO Double
-> DiscretePDF a
-> IO a
generateDiscrete01 :: forall a. IO Double -> DiscretePDF a -> IO a
generateDiscrete01 IO Double
g [] = forall a. HasCallStack => [Char] -> a
error [Char]
"Empty PDF: generateDiscrete01"
generateDiscrete01 IO Double
g [(a, Double)]
dpdf =
do Double
x <- IO Double
g
let loop :: Double -> [(a, Double)] -> a
loop Double
acc [(a
a, Double
p)] = a
a
loop Double
acc ((a
a, Double
p) : [(a, Double)]
dpdf) =
if Double
x forall a. Ord a => a -> a -> Bool
<= Double
acc forall a. Num a => a -> a -> a
+ Double
p
then a
a
else Double -> [(a, Double)] -> a
loop (Double
acc forall a. Num a => a -> a -> a
+ Double
p) [(a, Double)]
dpdf
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall {a}. Double -> [(a, Double)] -> a
loop Double
0 [(a, Double)]
dpdf
data GeneratorType = SimpleGenerator
| SimpleGeneratorWithSeed Word32
| CustomGenerator (IO Generator)
| CustomGenerator01 (IO Double)
newGenerator :: GeneratorType -> IO Generator
newGenerator :: GeneratorType -> IO Generator
newGenerator GeneratorType
tp =
case GeneratorType
tp of
GeneratorType
SimpleGenerator ->
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
Gen (PrimState m) -> m a
MWC.uniform forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO GenIO
MWC.createSystemRandom forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO Double -> IO Generator
newRandomGenerator01
SimpleGeneratorWithSeed Word32
x ->
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
Gen (PrimState m) -> m a
MWC.uniform forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) (v :: * -> *).
(PrimMonad m, Vector v Word32) =>
v Word32 -> m (Gen (PrimState m))
MWC.initialize (forall a. a -> Vector a
singleton Word32
x) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO Double -> IO Generator
newRandomGenerator01
CustomGenerator IO Generator
g ->
IO Generator
g
CustomGenerator01 IO Double
g ->
IO Double -> IO Generator
newRandomGenerator01 IO Double
g
newRandomGenerator :: RandomGen g => g -> IO Generator
newRandomGenerator :: forall g. RandomGen g => g -> IO Generator
newRandomGenerator g
g =
do IORef g
r <- forall a. a -> IO (IORef a)
newIORef g
g
let g1 :: IO Double
g1 = do g
g <- forall a. IORef a -> IO a
readIORef IORef g
r
let (Double
x, g
g') = forall a g. (Random a, RandomGen g) => g -> (a, g)
random g
g
forall a. IORef a -> a -> IO ()
writeIORef IORef g
r g
g'
forall (m :: * -> *) a. Monad m => a -> m a
return Double
x
IO Double -> IO Generator
newRandomGenerator01 IO Double
g1
newRandomGenerator01 :: IO Double -> IO Generator
newRandomGenerator01 :: IO Double -> IO Generator
newRandomGenerator01 IO Double
g =
do let g1 :: IO Double
g1 = IO Double
g
IO Double
g2 <- IO Double -> IO (IO Double)
newNormalGenerator01 IO Double
g1
let g3 :: Double -> Double -> IO Double
g3 Double
mu Double
nu = do { Double
x <- IO Double
g2; forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Double
mu forall a. Num a => a -> a -> a
+ Double
nu forall a. Num a => a -> a -> a
* Double
x }
g4 :: Double -> Double -> IO Double
g4 Double
mu Double
nu = do { Double
x <- IO Double
g2; forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Floating a => a -> a
exp (Double
mu forall a. Num a => a -> a -> a
+ Double
nu forall a. Num a => a -> a -> a
* Double
x) }
IORef Int
seqNoRef <- forall a. a -> IO (IORef a)
newIORef Int
0
let seqNo :: IO Int
seqNo = do { Int
x <- forall a. IORef a -> IO a
readIORef IORef Int
seqNoRef; forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
seqNoRef (forall a. Num a => a -> a -> a
+Int
1); forall (m :: * -> *) a. Monad m => a -> m a
return Int
x }
forall (m :: * -> *) a. Monad m => a -> m a
return Generator { generateUniform :: Double -> Double -> IO Double
generateUniform = IO Double -> Double -> Double -> IO Double
generateUniform01 IO Double
g1,
generateUniformInt :: Int -> Int -> IO Int
generateUniformInt = IO Double -> Int -> Int -> IO Int
generateUniformInt01 IO Double
g1,
generateTriangular :: Double -> Double -> Double -> IO Double
generateTriangular = IO Double -> Double -> Double -> Double -> IO Double
generateTriangular01 IO Double
g1,
generateNormal :: Double -> Double -> IO Double
generateNormal = Double -> Double -> IO Double
g3,
generateLogNormal :: Double -> Double -> IO Double
generateLogNormal = Double -> Double -> IO Double
g4,
generateExponential :: Double -> IO Double
generateExponential = IO Double -> Double -> IO Double
generateExponential01 IO Double
g1,
generateErlang :: Double -> Int -> IO Double
generateErlang = IO Double -> Double -> Int -> IO Double
generateErlang01 IO Double
g1,
generatePoisson :: Double -> IO Int
generatePoisson = IO Double -> Double -> IO Int
generatePoisson01 IO Double
g1,
generateBinomial :: Double -> Int -> IO Int
generateBinomial = IO Double -> Double -> Int -> IO Int
generateBinomial01 IO Double
g1,
generateGamma :: Double -> Double -> IO Double
generateGamma = IO Double -> IO Double -> Double -> Double -> IO Double
generateGamma01 IO Double
g2 IO Double
g1,
generateBeta :: Double -> Double -> IO Double
generateBeta = IO Double -> IO Double -> Double -> Double -> IO Double
generateBeta01 IO Double
g2 IO Double
g1,
generateWeibull :: Double -> Double -> IO Double
generateWeibull = IO Double -> Double -> Double -> IO Double
generateWeibull01 IO Double
g1,
generateDiscrete :: forall a. DiscretePDF a -> IO a
generateDiscrete = forall a. IO Double -> DiscretePDF a -> IO a
generateDiscrete01 IO Double
g1,
generateSequenceNo :: IO Int
generateSequenceNo = IO Int
seqNo }