{-# LANGUAGE FlexibleContexts #-}
module System.Random.MWC.CondensedTable (
CondensedTable
, CondensedTableV
, CondensedTableU
, genFromTable
, tableFromProbabilities
, tableFromWeights
, tableFromIntWeights
, tablePoisson
, tableBinomial
) where
import Control.Arrow (second,(***))
import Control.Monad.Primitive (PrimMonad(..))
import Data.Word
import Data.Int
import Data.Bits
import qualified Data.Vector.Generic as G
import Data.Vector.Generic ((++))
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector as V
import Data.Vector.Generic (Vector)
import Numeric.SpecFunctions (logFactorial)
import Prelude hiding ((++))
import System.Random.MWC
data CondensedTable v a =
CondensedTable
{-# UNPACK #-} !Word64 !(v a)
{-# UNPACK #-} !Word64 !(v a)
{-# UNPACK #-} !Word64 !(v a)
!(v a)
type CondensedTableU = CondensedTable U.Vector
type CondensedTableV = CondensedTable V.Vector
genFromTable :: (PrimMonad m, Vector v a) =>
CondensedTable v a -> Gen (PrimState m) -> m a
{-# INLINE genFromTable #-}
genFromTable :: CondensedTable v a -> Gen (PrimState m) -> m a
genFromTable table :: CondensedTable v a
table gen :: Gen (PrimState m)
gen = do
Word32
w <- Gen (PrimState m) -> m Word32
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
Gen (PrimState m) -> m a
uniform Gen (PrimState m)
gen
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
$ CondensedTable v a -> Word64 -> a
forall (v :: * -> *) a.
Vector v a =>
CondensedTable v a -> Word64 -> a
lookupTable CondensedTable v a
table (Word64 -> a) -> Word64 -> a
forall a b. (a -> b) -> a -> b
$ Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32
w :: Word32)
lookupTable :: Vector v a => CondensedTable v a -> Word64 -> a
{-# INLINE lookupTable #-}
lookupTable :: CondensedTable v a -> Word64 -> a
lookupTable (CondensedTable na :: Word64
na aa :: v a
aa nb :: Word64
nb bb :: v a
bb nc :: Word64
nc cc :: v a
cc dd :: v a
dd) i :: Word64
i
| Word64
i Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
na = v a
aa v a -> Word64 -> a
forall (v :: * -> *) a a. (Vector v a, Integral a) => v a -> a -> a
`at` ( Word64
i Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` 24)
| Word64
i Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
nb = v a
bb v a -> Word64 -> a
forall (v :: * -> *) a a. (Vector v a, Integral a) => v a -> a -> a
`at` ((Word64
i Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
na) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` 16)
| Word64
i Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
nc = v a
cc v a -> Word64 -> a
forall (v :: * -> *) a a. (Vector v a, Integral a) => v a -> a -> a
`at` ((Word64
i Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
nb) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` 8 )
| Bool
otherwise = v a
dd v a -> Word64 -> a
forall (v :: * -> *) a a. (Vector v a, Integral a) => v a -> a -> a
`at` ( Word64
i Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
nc)
where
at :: v a -> a -> a
at arr :: v a
arr j :: a
j = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
arr (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
j)
tableFromProbabilities
:: (Vector v (a,Word32), Vector v (a,Double), Vector v a, Vector v Word32)
=> v (a, Double) -> CondensedTable v a
{-# INLINE tableFromProbabilities #-}
tableFromProbabilities :: v (a, Double) -> CondensedTable v a
tableFromProbabilities v :: v (a, Double)
v
| v (a, Double) -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v (a, Double)
tbl = String -> String -> CondensedTable v a
forall a. String -> String -> a
pkgError "tableFromProbabilities" "empty vector of outcomes"
| Bool
otherwise = v (a, Word32) -> CondensedTable v a
forall (v :: * -> *) a.
(Vector v (a, Word32), Vector v a, Vector v Word32) =>
v (a, Word32) -> CondensedTable v a
tableFromIntWeights (v (a, Word32) -> CondensedTable v a)
-> v (a, Word32) -> CondensedTable v a
forall a b. (a -> b) -> a -> b
$ ((a, Double) -> (a, Word32)) -> v (a, Double) -> v (a, Word32)
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map ((Double -> Word32) -> (a, Double) -> (a, Word32)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ((Double -> Word32) -> (a, Double) -> (a, Word32))
-> (Double -> Word32) -> (a, Double) -> (a, Word32)
forall a b. (a -> b) -> a -> b
$ Double -> Word32
forall p. Integral p => Double -> p
toWeight (Double -> Word32) -> (Double -> Double) -> Double -> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
mlt)) v (a, Double)
tbl
where
mlt :: Double
mlt = 4.294967296e9
tbl :: v (a, Double)
tbl = ((a, Double) -> Bool) -> v (a, Double) -> v (a, Double)
forall (v :: * -> *) a. Vector v a => (a -> Bool) -> v a -> v a
G.filter ((Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> 0) (Double -> Bool) -> ((a, Double) -> Double) -> (a, Double) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, Double) -> Double
forall a b. (a, b) -> b
snd) v (a, Double)
v
toWeight :: Double -> p
toWeight w :: Double
w | Double
w Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
mlt Double -> Double -> Double
forall a. Num a => a -> a -> a
- 1 = 2p -> Int -> p
forall a b. (Num a, Integral b) => a -> b -> a
^(32::Int) p -> p -> p
forall a. Num a => a -> a -> a
- 1
| Bool
otherwise = Double -> p
forall a b. (RealFrac a, Integral b) => a -> b
round Double
w
tableFromWeights
:: (Vector v (a,Word32), Vector v (a,Double), Vector v a, Vector v Word32)
=> v (a, Double) -> CondensedTable v a
{-# INLINE tableFromWeights #-}
tableFromWeights :: v (a, Double) -> CondensedTable v a
tableFromWeights = v (a, Double) -> CondensedTable v a
forall (v :: * -> *) a.
(Vector v (a, Word32), Vector v (a, Double), Vector v a,
Vector v Word32) =>
v (a, Double) -> CondensedTable v a
tableFromProbabilities (v (a, Double) -> CondensedTable v a)
-> (v (a, Double) -> v (a, Double))
-> v (a, Double)
-> CondensedTable v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v (a, Double) -> v (a, Double)
forall (v :: * -> *) a b.
(Fractional b, Vector v (a, b)) =>
v (a, b) -> v (a, b)
normalize (v (a, Double) -> v (a, Double))
-> (v (a, Double) -> v (a, Double))
-> v (a, Double)
-> v (a, Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, Double) -> Bool) -> v (a, Double) -> v (a, Double)
forall (v :: * -> *) a. Vector v a => (a -> Bool) -> v a -> v a
G.filter ((Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> 0) (Double -> Bool) -> ((a, Double) -> Double) -> (a, Double) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, Double) -> Double
forall a b. (a, b) -> b
snd)
where
normalize :: v (a, b) -> v (a, b)
normalize v :: v (a, b)
v
| v (a, b) -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v (a, b)
v = String -> String -> v (a, b)
forall a. String -> String -> a
pkgError "tableFromWeights" "no positive weights"
| Bool
otherwise = ((a, b) -> (a, b)) -> v (a, b) -> v (a, b)
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map ((b -> b) -> (a, b) -> (a, b)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (b -> b -> b
forall a. Fractional a => a -> a -> a
/ b
s)) v (a, b)
v
where
s :: b
s = (b -> (a, b) -> b) -> b -> v (a, b) -> b
forall (v :: * -> *) b a.
Vector v b =>
(a -> b -> a) -> a -> v b -> a
G.foldl' (((a, b) -> b -> b) -> b -> (a, b) -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip (((a, b) -> b -> b) -> b -> (a, b) -> b)
-> ((a, b) -> b -> b) -> b -> (a, b) -> b
forall a b. (a -> b) -> a -> b
$ b -> b -> b
forall a. Num a => a -> a -> a
(+) (b -> b -> b) -> ((a, b) -> b) -> (a, b) -> b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, b) -> b
forall a b. (a, b) -> b
snd) 0 v (a, b)
v
tableFromIntWeights :: (Vector v (a,Word32), Vector v a, Vector v Word32)
=> v (a, Word32)
-> CondensedTable v a
{-# INLINE tableFromIntWeights #-}
tableFromIntWeights :: v (a, Word32) -> CondensedTable v a
tableFromIntWeights v :: v (a, Word32)
v
| Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 0 = String -> String -> CondensedTable v a
forall a. String -> String -> a
pkgError "tableFromIntWeights" "empty table"
| Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 1 = let m :: Word64
m = 2Word64 -> Int -> Word64
forall a b. (Num a, Integral b) => a -> b -> a
^(32::Int) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- 1
in Word64
-> v a
-> Word64
-> v a
-> Word64
-> v a
-> v a
-> CondensedTable v a
forall (v :: * -> *) a.
Word64
-> v a
-> Word64
-> v a
-> Word64
-> v a
-> v a
-> CondensedTable v a
CondensedTable
Word64
m (Int -> a -> v a
forall (v :: * -> *) a. Vector v a => Int -> a -> v a
G.replicate 256 (a -> v a) -> a -> v a
forall a b. (a -> b) -> a -> b
$ (a, Word32) -> a
forall a b. (a, b) -> a
fst ((a, Word32) -> a) -> (a, Word32) -> a
forall a b. (a -> b) -> a -> b
$ v (a, Word32) -> (a, Word32)
forall (v :: * -> *) a. Vector v a => v a -> a
G.head v (a, Word32)
tbl)
Word64
m v a
forall (v :: * -> *) a. Vector v a => v a
G.empty
Word64
m v a
forall (v :: * -> *) a. Vector v a => v a
G.empty
v a
forall (v :: * -> *) a. Vector v a => v a
G.empty
| Bool
otherwise = Word64
-> v a
-> Word64
-> v a
-> Word64
-> v a
-> v a
-> CondensedTable v a
forall (v :: * -> *) a.
Word64
-> v a
-> Word64
-> v a
-> Word64
-> v a
-> v a
-> CondensedTable v a
CondensedTable
Word64
na v a
aa
Word64
nb v a
bb
Word64
nc v a
cc
v a
dd
where
tbl :: v (a, Word32)
tbl = ((a, Word32) -> Bool) -> v (a, Word32) -> v (a, Word32)
forall (v :: * -> *) a. Vector v a => (a -> Bool) -> v a -> v a
G.filter ((Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
/=0) (Word32 -> Bool) -> ((a, Word32) -> Word32) -> (a, Word32) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, Word32) -> Word32
forall a b. (a, b) -> b
snd) v (a, Word32)
v
n :: Int
n = v (a, Word32) -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (a, Word32)
tbl
table :: v (a, Word32)
table = (v a -> v Word32 -> v (a, Word32))
-> (v a, v Word32) -> v (a, Word32)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry v a -> v Word32 -> v (a, Word32)
forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v (a, b)) =>
v a -> v b -> v (a, b)
G.zip ((v a, v Word32) -> v (a, Word32))
-> (v a, v Word32) -> v (a, Word32)
forall a b. (a -> b) -> a -> b
$ v a -> v a
forall a. a -> a
id (v a -> v a)
-> (v Word32 -> v Word32) -> (v a, v Word32) -> (v a, v Word32)
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** v Word32 -> v Word32
forall (v :: * -> *). Vector v Word32 => v Word32 -> v Word32
correctWeights ((v a, v Word32) -> (v a, v Word32))
-> (v a, v Word32) -> (v a, v Word32)
forall a b. (a -> b) -> a -> b
$ v (a, Word32) -> (v a, v Word32)
forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v (a, b)) =>
v (a, b) -> (v a, v b)
G.unzip v (a, Word32)
tbl
mkTable :: Int -> v a
mkTable d :: Int
d =
((a, Word32) -> v a) -> v (a, Word32) -> v a
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> v b) -> v a -> v b
G.concatMap (\(x :: a
x,w :: Word32
w) -> Int -> a -> v a
forall (v :: * -> *) a. Vector v a => Int -> a -> v a
G.replicate (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> Int) -> Word32 -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Word32 -> Word32
digit Int
d Word32
w) a
x) v (a, Word32)
table
len :: v a -> Word64
len = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64) -> (v a -> Int) -> v a -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length
aa :: v a
aa = Int -> v a
mkTable 0
bb :: v a
bb = Int -> v a
mkTable 1
cc :: v a
cc = Int -> v a
mkTable 2
dd :: v a
dd = Int -> v a
mkTable 3
na :: Word64
na = v a -> Word64
len v a
aa Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` 24
nb :: Word64
nb = Word64
na Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ (v a -> Word64
len v a
bb Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` 16)
nc :: Word64
nc = Word64
nb Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ (v a -> Word64
len v a
cc Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` 8)
digit :: Int -> Word32 -> Word32
digit :: Int -> Word32 -> Word32
digit 0 x :: Word32
x = Word32
x Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftR` 24
digit 1 x :: Word32
x = (Word32
x Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftR` 16) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. 0xff
digit 2 x :: Word32
x = (Word32
x Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftR` 8 ) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. 0xff
digit 3 x :: Word32
x = Word32
x Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. 0xff
digit _ _ = String -> String -> Word32
forall a. String -> String -> a
pkgError "digit" "the impossible happened!?"
{-# INLINE digit #-}
correctWeights :: G.Vector v Word32 => v Word32 -> v Word32
{-# INLINE correctWeights #-}
correctWeights :: v Word32 -> v Word32
correctWeights v :: v Word32
v = (forall s. ST s (Mutable v s Word32)) -> v Word32
forall (v :: * -> *) a.
Vector v a =>
(forall s. ST s (Mutable v s a)) -> v a
G.create ((forall s. ST s (Mutable v s Word32)) -> v Word32)
-> (forall s. ST s (Mutable v s Word32)) -> v Word32
forall a b. (a -> b) -> a -> b
$ do
let
s :: Int64
s = (Int64 -> Word32 -> Int64) -> Int64 -> v Word32 -> Int64
forall (v :: * -> *) b a.
Vector v b =>
(a -> b -> a) -> a -> v b -> a
G.foldl' ((Word32 -> Int64 -> Int64) -> Int64 -> Word32 -> Int64
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Word32 -> Int64 -> Int64) -> Int64 -> Word32 -> Int64)
-> (Word32 -> Int64 -> Int64) -> Int64 -> Word32 -> Int64
forall a b. (a -> b) -> a -> b
$ Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
(+) (Int64 -> Int64 -> Int64)
-> (Word32 -> Int64) -> Word32 -> Int64 -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word32 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral) 0 v Word32
v :: Int64
n :: Int
n = v Word32 -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v Word32
v
Mutable v s Word32
arr <- v Word32 -> ST s (Mutable v (PrimState (ST s)) Word32)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
G.thaw v Word32
v
let loop :: Word32 -> Int -> a -> ST s ()
loop lim :: Word32
lim i :: Int
i delta :: a
delta
| a
delta a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== 0 = () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Word32 -> Int -> a -> ST s ()
loop 1 0 a
delta
| Bool
otherwise = do
Word32
w <- Mutable v (PrimState (ST s)) Word32 -> Int -> ST s Word32
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
M.read Mutable v s Word32
Mutable v (PrimState (ST s)) Word32
arr Int
i
case () of
_| Word32
w Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
< Word32
lim -> Word32 -> Int -> a -> ST s ()
loop Word32
lim (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) a
delta
| a
delta a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< 0 -> Mutable v (PrimState (ST s)) Word32 -> Int -> Word32 -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
M.write Mutable v s Word32
Mutable v (PrimState (ST s)) Word32
arr Int
i (Word32
w Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ 1) ST s () -> ST s () -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word32 -> Int -> a -> ST s ()
loop Word32
lim (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) (a
delta a -> a -> a
forall a. Num a => a -> a -> a
+ 1)
| Bool
otherwise -> Mutable v (PrimState (ST s)) Word32 -> Int -> Word32 -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
M.write Mutable v s Word32
Mutable v (PrimState (ST s)) Word32
arr Int
i (Word32
w Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- 1) ST s () -> ST s () -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word32 -> Int -> a -> ST s ()
loop Word32
lim (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) (a
delta a -> a -> a
forall a. Num a => a -> a -> a
- 1)
Word32 -> Int -> Int64 -> ST s ()
forall a. (Num a, Ord a) => Word32 -> Int -> a -> ST s ()
loop 255 0 (Int64
s Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- 2Int64 -> Int -> Int64
forall a b. (Num a, Integral b) => a -> b -> a
^(32::Int))
Mutable v s Word32 -> ST s (Mutable v s Word32)
forall (m :: * -> *) a. Monad m => a -> m a
return Mutable v s Word32
arr
tablePoisson :: Double -> CondensedTableU Int
tablePoisson :: Double -> CondensedTableU Int
tablePoisson = Vector (Int, Double) -> CondensedTableU Int
forall (v :: * -> *) a.
(Vector v (a, Word32), Vector v (a, Double), Vector v a,
Vector v Word32) =>
v (a, Double) -> CondensedTable v a
tableFromProbabilities (Vector (Int, Double) -> CondensedTableU Int)
-> (Double -> Vector (Int, Double))
-> Double
-> CondensedTableU Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Vector (Int, Double)
make
where
make :: Double -> Vector (Int, Double)
make lam :: Double
lam
| Double
lam Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< 0 = String -> String -> Vector (Int, Double)
forall a. String -> String -> a
pkgError "tablePoisson" "negative lambda"
| Double
lam Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< 22.8 = ((Double, Int) -> Maybe ((Int, Double), (Double, Int)))
-> (Double, Int) -> Vector (Int, Double)
forall a b. Unbox a => (b -> Maybe (a, b)) -> b -> Vector a
U.unfoldr (Double, Int) -> Maybe ((Int, Double), (Double, Int))
forall b.
Integral b =>
(Double, b) -> Maybe ((b, Double), (Double, b))
unfoldForward (Double -> Double
forall a. Floating a => a -> a
exp (-Double
lam), 0)
| Bool
otherwise = ((Double, Int) -> Maybe ((Int, Double), (Double, Int)))
-> (Double, Int) -> Vector (Int, Double)
forall a b. Unbox a => (b -> Maybe (a, b)) -> b -> Vector a
U.unfoldr (Double, Int) -> Maybe ((Int, Double), (Double, Int))
forall b.
Integral b =>
(Double, b) -> Maybe ((b, Double), (Double, b))
unfoldForward (Double
pMax, Int
nMax)
Vector (Int, Double)
-> Vector (Int, Double) -> Vector (Int, Double)
forall (v :: * -> *) a. Vector v a => v a -> v a -> v a
++ Vector (Int, Double) -> Vector (Int, Double)
forall a. Unbox a => Vector a -> Vector a
U.tail (((Double, Int) -> Maybe ((Int, Double), (Double, Int)))
-> (Double, Int) -> Vector (Int, Double)
forall a b. Unbox a => (b -> Maybe (a, b)) -> b -> Vector a
U.unfoldr (Double, Int) -> Maybe ((Int, Double), (Double, Int))
forall b.
Integral b =>
(Double, b) -> Maybe ((b, Double), (Double, b))
unfoldBackward (Double
pMax, Int
nMax))
where
nMax :: Int
nMax = Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
lam :: Int
pMax :: Double
pMax = Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
nMax Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
lam Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lam Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
forall a. Integral a => a -> Double
logFactorial Int
nMax
unfoldForward :: (Double, b) -> Maybe ((b, Double), (Double, b))
unfoldForward (p :: Double
p,i :: b
i)
| Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
minP = Maybe ((b, Double), (Double, b))
forall a. Maybe a
Nothing
| Bool
otherwise = ((b, Double), (Double, b)) -> Maybe ((b, Double), (Double, b))
forall a. a -> Maybe a
Just ( (b
i,Double
p)
, (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
lam Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ b -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (b
ib -> b -> b
forall a. Num a => a -> a -> a
+1), b
ib -> b -> b
forall a. Num a => a -> a -> a
+1)
)
unfoldBackward :: (Double, b) -> Maybe ((b, Double), (Double, b))
unfoldBackward (p :: Double
p,i :: b
i)
| Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
minP = Maybe ((b, Double), (Double, b))
forall a. Maybe a
Nothing
| Bool
otherwise = ((b, Double), (Double, b)) -> Maybe ((b, Double), (Double, b))
forall a. a -> Maybe a
Just ( (b
i,Double
p)
, (Double
p Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
lam Double -> Double -> Double
forall a. Num a => a -> a -> a
* b -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral b
i, b
ib -> b -> b
forall a. Num a => a -> a -> a
-1)
)
minP :: Double
minP = 1.1641532182693481e-10
tableBinomial :: Int
-> Double
-> CondensedTableU Int
tableBinomial :: Int -> Double -> CondensedTableU Int
tableBinomial n :: Int
n p :: Double
p = Vector (Int, Double) -> CondensedTableU Int
forall (v :: * -> *) a.
(Vector v (a, Word32), Vector v (a, Double), Vector v a,
Vector v Word32) =>
v (a, Double) -> CondensedTable v a
tableFromProbabilities Vector (Int, Double)
makeBinom
where
makeBinom :: Vector (Int, Double)
makeBinom
| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= 0 = String -> String -> Vector (Int, Double)
forall a. String -> String -> a
pkgError "tableBinomial" "non-positive number of tries"
| Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== 0 = (Int, Double) -> Vector (Int, Double)
forall a. Unbox a => a -> Vector a
U.singleton (0,1)
| Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== 1 = (Int, Double) -> Vector (Int, Double)
forall a. Unbox a => a -> Vector a
U.singleton (Int
n,1)
| Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> 0 Bool -> Bool -> Bool
&& Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< 1 = Int
-> ((Double, Int) -> Maybe ((Int, Double), (Double, Int)))
-> (Double, Int)
-> Vector (Int, Double)
forall a b. Unbox a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
U.unfoldrN (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) (Double, Int) -> Maybe ((Int, Double), (Double, Int))
unfolder ((1Double -> Double -> Double
forall a. Num a => a -> a -> a
-Double
p)Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^Int
n, 0)
| Bool
otherwise = String -> String -> Vector (Int, Double)
forall a. String -> String -> a
pkgError "tableBinomial" "probability is out of range"
where
h :: Double
h = Double
p Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p)
unfolder :: (Double, Int) -> Maybe ((Int, Double), (Double, Int))
unfolder (t :: Double
t,i :: Int
i) = ((Int, Double), (Double, Int))
-> Maybe ((Int, Double), (Double, Int))
forall a. a -> Maybe a
Just ( (Int
i,Double
t)
, (Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Double) -> Int -> Double
forall a b. (a -> b) -> a -> b
$ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i1) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i1, Int
i1) )
where i1 :: Int
i1 = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1
pkgError :: String -> String -> a
pkgError :: String -> String -> a
pkgError func :: String
func err :: String
err =
String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> ([String] -> String) -> [String] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([String] -> a) -> [String] -> a
forall a b. (a -> b) -> a -> b
$ ["System.Random.MWC.CondensedTable.", String
func, ": ", String
err]