{-# LANGUAGE BangPatterns #-}

-- SPDX-License-Identifier: MPL-2.0

{- | Explicit, immutable random streams. State format @hmath-rng-1@ is decimal
'Show' syntax and is independent of host byte order. A split consumes two
words and hashes each to a new seed; it promises deterministic distinct seeds,
not disjoint mathematical subsequences. Generators are for simulation; even
ChaCha's 64-bit convenience seed is not a cryptographic key interface.

PCG32 XSH-RR: https://www.pcg-random.org/using-pcg-c-basic.html
xoshiro/xoroshiro: https://prng.di.unimi.it/
Counter algorithms: https://www.thesalmons.org/john/random123/
ChaCha20 rounds and counter layout: https://www.rfc-editor.org/rfc/rfc8439
-}
module HMath.Random (
    Generator (..),
    RNG,
    seedRNG,
    nextWord64,
    cloneRNG,
    splitRNG,
    serializeRNG,
    restoreRNG,
    uniform,
    uniformOpen,
    uniformInt,
    sampleN,
    normal,
    boxMuller,
    marsagliaNormal,
    zigguratNormal,
    exponential,
    gamma,
    beta,
    poisson,
    binomial,
    weighted,
    AliasTable,
    aliasTable,
    aliasSample,
    reservoir,
    shuffle,
    withReplacement,
    withoutReplacement,
    stratified,
    importance,
    rejection,
    inverseTransform,
    metropolisHastings,
    gibbs,
    hamiltonianMonteCarlo,
    monteCarlo,
    halton,
    sobol,
    faure,
    latinHypercube,
    randomQuality,
    entropySeed,
) where

import Data.Bits
import qualified Data.Foldable as Foldable (foldl')
import Data.Int (Int32)
import Data.List (transpose)
import Data.Word
import HMath.Error
import System.IO (IOMode (ReadMode), hGetChar, withBinaryFile)
import Text.Read (readMaybe)

-- | Algorithm variant. PCG32 concatenates two XSH-RR outputs, MT19937 two 32-bit outputs; counter engines pack adjacent words little endian.
data Generator
    = PCG32
    | Xoshiro256StarStar
    | Xoroshiro128Plus
    | MT19937
    | Philox4x32
    | Threefry2x64
    | ChaCha20
    deriving (Generator -> Generator -> Bool
(Generator -> Generator -> Bool)
-> (Generator -> Generator -> Bool) -> Eq Generator
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Generator -> Generator -> Bool
== :: Generator -> Generator -> Bool
$c/= :: Generator -> Generator -> Bool
/= :: Generator -> Generator -> Bool
Eq, Int -> Generator -> ShowS
[Generator] -> ShowS
Generator -> String
(Int -> Generator -> ShowS)
-> (Generator -> String)
-> ([Generator] -> ShowS)
-> Show Generator
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Generator -> ShowS
showsPrec :: Int -> Generator -> ShowS
$cshow :: Generator -> String
show :: Generator -> String
$cshowList :: [Generator] -> ShowS
showList :: [Generator] -> ShowS
Show, ReadPrec [Generator]
ReadPrec Generator
Int -> ReadS Generator
ReadS [Generator]
(Int -> ReadS Generator)
-> ReadS [Generator]
-> ReadPrec Generator
-> ReadPrec [Generator]
-> Read Generator
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS Generator
readsPrec :: Int -> ReadS Generator
$creadList :: ReadS [Generator]
readList :: ReadS [Generator]
$creadPrec :: ReadPrec Generator
readPrec :: ReadPrec Generator
$creadListPrec :: ReadPrec [Generator]
readListPrec :: ReadPrec [Generator]
Read, Int -> Generator
Generator -> Int
Generator -> [Generator]
Generator -> Generator
Generator -> Generator -> [Generator]
Generator -> Generator -> Generator -> [Generator]
(Generator -> Generator)
-> (Generator -> Generator)
-> (Int -> Generator)
-> (Generator -> Int)
-> (Generator -> [Generator])
-> (Generator -> Generator -> [Generator])
-> (Generator -> Generator -> [Generator])
-> (Generator -> Generator -> Generator -> [Generator])
-> Enum Generator
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: Generator -> Generator
succ :: Generator -> Generator
$cpred :: Generator -> Generator
pred :: Generator -> Generator
$ctoEnum :: Int -> Generator
toEnum :: Int -> Generator
$cfromEnum :: Generator -> Int
fromEnum :: Generator -> Int
$cenumFrom :: Generator -> [Generator]
enumFrom :: Generator -> [Generator]
$cenumFromThen :: Generator -> Generator -> [Generator]
enumFromThen :: Generator -> Generator -> [Generator]
$cenumFromTo :: Generator -> Generator -> [Generator]
enumFromTo :: Generator -> Generator -> [Generator]
$cenumFromThenTo :: Generator -> Generator -> Generator -> [Generator]
enumFromThenTo :: Generator -> Generator -> Generator -> [Generator]
Enum, Generator
Generator -> Generator -> Bounded Generator
forall a. a -> a -> Bounded a
$cminBound :: Generator
minBound :: Generator
$cmaxBound :: Generator
maxBound :: Generator
Bounded)

-- | Opaque immutable generator state, including any buffered block outputs. Equality includes the complete continuation state.
data RNG = RNG !Generator ![Word64] ![Word64] deriving (RNG -> RNG -> Bool
(RNG -> RNG -> Bool) -> (RNG -> RNG -> Bool) -> Eq RNG
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RNG -> RNG -> Bool
== :: RNG -> RNG -> Bool
$c/= :: RNG -> RNG -> Bool
/= :: RNG -> RNG -> Bool
Eq, Int -> RNG -> ShowS
[RNG] -> ShowS
RNG -> String
(Int -> RNG -> ShowS)
-> (RNG -> String) -> ([RNG] -> ShowS) -> Show RNG
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RNG -> ShowS
showsPrec :: Int -> RNG -> ShowS
$cshow :: RNG -> String
show :: RNG -> String
$cshowList :: [RNG] -> ShowS
showList :: [RNG] -> ShowS
Show)

-- Parsing uses a private raw type so Read cannot bypass the opaque invariant.
data RawRNG = RawRNG Generator [Word64] [Word64] deriving (ReadPrec [RawRNG]
ReadPrec RawRNG
Int -> ReadS RawRNG
ReadS [RawRNG]
(Int -> ReadS RawRNG)
-> ReadS [RawRNG]
-> ReadPrec RawRNG
-> ReadPrec [RawRNG]
-> Read RawRNG
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS RawRNG
readsPrec :: Int -> ReadS RawRNG
$creadList :: ReadS [RawRNG]
readList :: ReadS [RawRNG]
$creadPrec :: ReadPrec RawRNG
readPrec :: ReadPrec RawRNG
$creadListPrec :: ReadPrec [RawRNG]
readListPrec :: ReadPrec [RawRNG]
Read, Int -> RawRNG -> ShowS
[RawRNG] -> ShowS
RawRNG -> String
(Int -> RawRNG -> ShowS)
-> (RawRNG -> String) -> ([RawRNG] -> ShowS) -> Show RawRNG
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RawRNG -> ShowS
showsPrec :: Int -> RawRNG -> ShowS
$cshow :: RawRNG -> String
show :: RawRNG -> String
$cshowList :: [RawRNG] -> ShowS
showList :: [RawRNG] -> ShowS
Show)

-- | Copy an immutable state; the copy produces exactly the same future outputs.
cloneRNG :: RNG -> RNG
cloneRNG :: RNG -> RNG
cloneRNG = RNG -> RNG
forall a. a -> a
id

mix :: Word64 -> Word64
mix :: Word64 -> Word64
mix Word64
x =
    let a :: Word64
a = (Word64
x Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
x Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
30)) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
0xbf58476d1ce4e5b9
        b :: Word64
b = (Word64
a Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
a Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
27)) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
0x94d049bb133111eb
     in Word64
b Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
b Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
31)

-- | Deterministically encode a 64-bit seed. Shift engines use SplitMix64 expansion; MT uses the low 32 bits and its standard initialization recurrence. PCG uses a fixed odd stream increment and an offset initial state. Counter engines use the seed as key; ChaCha expands it into a 256-bit simulation key.
seedRNG :: Generator -> Word64 -> RNG
seedRNG :: Generator -> Word64 -> RNG
seedRNG Generator
g Word64
s = Generator -> [Word64] -> [Word64] -> RNG
RNG Generator
g [Word64]
state []
  where
    seeds :: [Word64]
seeds = (Word64 -> Word64) -> [Word64] -> [Word64]
forall a b. (a -> b) -> [a] -> [b]
map (Word64 -> Word64
mix (Word64 -> Word64) -> (Word64 -> Word64) -> Word64 -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word64
s Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+) (Word64 -> Word64) -> (Word64 -> Word64) -> Word64 -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
0x9e3779b97f4a7c15)) [Word64
1 ..]
    state :: [Word64]
state = case Generator
g of
        Generator
PCG32 -> [Word64
s Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
0x853c49e6748fea9b, Word64
0xda3e39cb94b95bdb]
        Generator
Xoshiro256StarStar -> Int -> [Word64] -> [Word64]
forall a. Int -> [a] -> [a]
take Int
4 [Word64]
seeds
        Generator
Xoroshiro128Plus -> Int -> [Word64] -> [Word64]
forall a. Int -> [a] -> [a]
take Int
2 [Word64]
seeds
        Generator
MT19937 -> Int -> [Word64] -> [Word64]
forall a. Int -> [a] -> [a]
take Int
624 ((Word32 -> Word64) -> [Word32] -> [Word64]
forall a b. (a -> b) -> [a] -> [b]
map Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> [Word32]
iterateMT (Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
s))) [Word64] -> [Word64] -> [Word64]
forall a. [a] -> [a] -> [a]
++ [Word64
624]
        Generator
Philox4x32 -> [Word64
0, Word64
s]
        Generator
Threefry2x64 -> [Word64
0, Word64
s]
        Generator
ChaCha20 -> [Word64
0] [Word64] -> [Word64] -> [Word64]
forall a. [a] -> [a] -> [a]
++ Int -> [Word64] -> [Word64]
forall a. Int -> [a] -> [a]
take Int
4 [Word64]
seeds
    iterateMT :: Word32 -> [Word32]
    iterateMT :: Word32 -> [Word32]
iterateMT Word32
a = Word32
a Word32 -> [Word32] -> [Word32]
forall a. a -> [a] -> [a]
: Word32 -> Word32 -> [Word32]
go Word32
1 Word32
a
    go :: Word32 -> Word32 -> [Word32]
    go :: Word32 -> Word32 -> [Word32]
go Word32
i Word32
a = let b :: Word32
b = Word32
1812433253 Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
* (Word32
a Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
`xor` (Word32
a Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftR` Int
30)) Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
i in Word32
b Word32 -> [Word32] -> [Word32]
forall a. a -> [a] -> [a]
: Word32 -> Word32 -> [Word32]
go (Word32
i Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
1) Word32
b

-- | Produce one word and its successor state. Fixed-width arithmetic wraps modulo 2^32 or 2^64 as required by the algorithm.
nextWord64 :: RNG -> (Word64, RNG)
nextWord64 :: RNG -> (Word64, RNG)
nextWord64 (RNG Generator
g [Word64]
s (Word64
w : [Word64]
ws)) = (Word64
w, Generator -> [Word64] -> [Word64] -> RNG
RNG Generator
g [Word64]
s [Word64]
ws)
nextWord64 (RNG Generator
g [Word64]
s []) = case (Generator
g, [Word64]
s) of
    (Generator
PCG32, [Word64
x, Word64
inc]) ->
        let (Word32
a, Word64
y) = Word64 -> Word64 -> (Word32, Word64)
pcg Word64
x Word64
inc; (Word32
b, Word64
z) = Word64 -> Word64 -> (Word32, Word64)
pcg Word64
y Word64
inc
         in (Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
a Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
32 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
b, Generator -> [Word64] -> [Word64] -> RNG
RNG Generator
g [Word64
z, Word64
inc] [])
    (Generator
Xoshiro256StarStar, [Word64
a, Word64
b, Word64
c, Word64
d]) ->
        let t :: Word64
t = Word64
b Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
17
            c1 :: Word64
c1 = Word64
c Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
a
            d1 :: Word64
d1 = Word64
d Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
b
            b1 :: Word64
b1 = Word64
b Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
c1
            a1 :: Word64
a1 = Word64
a Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
d1
            c2 :: Word64
c2 = Word64
c1 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
t
            d2 :: Word64
d2 = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
rotateL Word64
d1 Int
45
         in (Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
rotateL (Word64
b Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
5) Int
7 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
9, Generator -> [Word64] -> [Word64] -> RNG
RNG Generator
g [Word64
a1, Word64
b1, Word64
c2, Word64
d2] [])
    (Generator
Xoroshiro128Plus, [Word64
a, Word64
b]) ->
        let t :: Word64
t = Word64
b Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
a in (Word64
a Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
b, Generator -> [Word64] -> [Word64] -> RNG
RNG Generator
g [Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
rotateL Word64
a Int
24 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
t Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
t Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
16), Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
rotateL Word64
t Int
37] [])
    (Generator
MT19937, [Word64]
_) ->
        let (Word32
a, [Word64]
s1) = [Word64] -> (Word32, [Word64])
mtNext [Word64]
s; (Word32
b, [Word64]
s2) = [Word64] -> (Word32, [Word64])
mtNext [Word64]
s1
         in (Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
a Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
32 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
b, Generator -> [Word64] -> [Word64] -> RNG
RNG Generator
g [Word64]
s2 [])
    (Generator
Philox4x32, [Word64
c, Word64
k]) -> let ws :: [Word32]
ws = Word64 -> Word64 -> [Word32]
philox Word64
c Word64
k in RNG -> (Word64, RNG)
nextWord64 (Generator -> [Word64] -> [Word64] -> RNG
RNG Generator
g [Word64
c Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
1, Word64
k] ([Word32] -> [Word64]
forall {a} {a}. (Bits a, Integral a, Num a) => [a] -> [a]
pack [Word32]
ws))
    (Generator
Threefry2x64, [Word64
c, Word64
k]) -> let (Word64
a, Word64
b) = Word64 -> Word64 -> (Word64, Word64)
threefry Word64
c Word64
k in (Word64
a, Generator -> [Word64] -> [Word64] -> RNG
RNG Generator
g [Word64
c Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
1, Word64
k] [Word64
b])
    (Generator
ChaCha20, Word64
c : [Word64]
ks) -> RNG -> (Word64, RNG)
nextWord64 (Generator -> [Word64] -> [Word64] -> RNG
RNG Generator
g ((Word64
c Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
1) Word64 -> [Word64] -> [Word64]
forall a. a -> [a] -> [a]
: [Word64]
ks) ([Word32] -> [Word64]
forall {a} {a}. (Bits a, Integral a, Num a) => [a] -> [a]
pack (Word64 -> [Word64] -> [Word32]
chacha Word64
c [Word64]
ks)))
    (Generator, [Word64])
_ -> RNG -> (Word64, RNG)
nextWord64 (Generator -> Word64 -> RNG
seedRNG Generator
g Word64
0) -- unreachable through the opaque, validated interface
  where
    pcg :: Word64 -> Word64 -> (Word32, Word64)
    pcg :: Word64 -> Word64 -> (Word32, Word64)
pcg Word64
x Word64
inc =
        ( Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
rotateR (Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (((Word64
x Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
18) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
x) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
27)) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
x Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
59))
        , Word64
x Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
6364136223846793005 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ (Word64
inc Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64
1)
        )
    pack :: [a] -> [a]
pack [] = []
    pack (a
a : a
b : [a]
xs) = (a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
a a -> a -> a
forall a. Bits a => a -> a -> a
.|. a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
b a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
32) a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a]
pack [a]
xs
    pack [a]
_ = []

mtNext :: [Word64] -> (Word32, [Word64])
mtNext :: [Word64] -> (Word32, [Word64])
mtNext [Word64]
s = (Word32 -> Word32
forall {a}. (Bits a, Num a) => a -> a
temper (Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Word64]
xs [Word64] -> Int -> Word64
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)), Int -> [Word64] -> [Word64]
forall a. Int -> [a] -> [a]
take Int
624 [Word64]
xs [Word64] -> [Word64] -> [Word64]
forall a. [a] -> [a] -> [a]
++ [Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)])
  where
    i0 :: Int
i0 = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Word64] -> Word64
forall a. HasCallStack => [a] -> a
last [Word64]
s) :: Int
    old :: [Word32]
old = (Word64 -> Word32) -> [Word64] -> [Word32]
forall a b. (a -> b) -> [a] -> [b]
map Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> [Word64] -> [Word64]
forall a. Int -> [a] -> [a]
take Int
624 [Word64]
s) :: [Word32]
    twist :: [Word32]
twist = ([Word32] -> Int -> [Word32]) -> [Word32] -> [Int] -> [Word32]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' [Word32] -> Int -> [Word32]
forall {a}. (Bits a, Integral a) => [a] -> Int -> [a]
upd [Word32]
old [Int
0 .. Int
623]
    upd :: [a] -> Int -> [a]
upd [a]
a Int
j =
        let y :: a
y = ([a]
a [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
j a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
0x80000000) a -> a -> a
forall a. Bits a => a -> a -> a
.|. ([a]
a [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! ((Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
624) a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
0x7fffffff)
            v :: a
v = [a]
a [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! ((Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
397) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
624) a -> a -> a
forall a. Bits a => a -> a -> a
`xor` (a
y a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
1) a -> a -> a
forall a. Bits a => a -> a -> a
`xor` (if a -> Bool
forall a. Integral a => a -> Bool
odd a
y then a
0x9908b0df else a
0)
         in Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
j [a]
a [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
v] [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
a
    xs :: [Word64]
xs = if Int
i0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
624 then (Word32 -> Word64) -> [Word32] -> [Word64]
forall a b. (a -> b) -> [a] -> [b]
map Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Word32]
twist else Int -> [Word64] -> [Word64]
forall a. Int -> [a] -> [a]
take Int
624 [Word64]
s
    i :: Int
i = if Int
i0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
624 then Int
0 else Int
i0
    temper :: a -> a
temper a
x =
        let a :: a
a = a
x a -> a -> a
forall a. Bits a => a -> a -> a
`xor` (a
x a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
11)
            b :: a
b = a
a a -> a -> a
forall a. Bits a => a -> a -> a
`xor` ((a
a a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
7) a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
0x9d2c5680)
            c :: a
c = a
b a -> a -> a
forall a. Bits a => a -> a -> a
`xor` ((a
b a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
15) a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
0xefc60000)
         in a
c a -> a -> a
forall a. Bits a => a -> a -> a
`xor` (a
c a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
18)
philox :: Word64 -> Word64 -> [Word32]
philox :: Word64 -> Word64 -> [Word32]
philox Word64
c Word64
k =
    Int -> [Word32] -> Word32 -> Word32 -> [Word32]
go
        Int
10
        [Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
c, Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
c Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
32), Word32
0, Word32
0]
        (Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
k)
        (Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
k Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
32))
  where
    go :: Int -> [Word32] -> Word32 -> Word32 -> [Word32]
    go :: Int -> [Word32] -> Word32 -> Word32 -> [Word32]
go Int
0 [Word32]
x Word32
_ Word32
_ = [Word32]
x
    go Int
n [Word32
a, Word32
b, Word32
d, Word32
e] Word32
p Word32
q =
        let x :: Word64
x = Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
a Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
0xd2511f53 :: Word64; y :: Word64
y = Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
d Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
0xcd9e8d57 :: Word64
         in Int -> [Word32] -> Word32 -> Word32 -> [Word32]
go
                (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                [ Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
y Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
32) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
`xor` Word32
b Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
`xor` Word32
p
                , Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
y
                , Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
x Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
32) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
`xor` Word32
e Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
`xor` Word32
q
                , Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
x
                ]
                (Word32
p Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
0x9e3779b9)
                (Word32
q Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
0xbb67ae85)
    go Int
_ [Word32]
x Word32
_ Word32
_ = [Word32]
x
threefry :: Word64 -> Word64 -> (Word64, Word64)
threefry :: Word64 -> Word64 -> (Word64, Word64)
threefry Word64
c Word64
k = Int -> Word64 -> Word64 -> (Word64, Word64)
go Int
0 (Word64
c Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
k) Word64
0
  where
    keys :: [Word64]
keys = [Word64
k, Word64
0, Word64
0x1bd11bdaa9fc1a22 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
k]
    rots :: [Int]
rots = [Int
16, Int
42, Int
12, Int
31, Int
16, Int
32, Int
24, Int
21]
    go :: Int -> Word64 -> Word64 -> (Word64, Word64)
    go :: Int -> Word64 -> Word64 -> (Word64, Word64)
go Int
20 Word64
a Word64
b = (Word64
a, Word64
b)
    go Int
n Word64
a Word64
b =
        let x :: Word64
x = Word64
a Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
b
            y :: Word64
y = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
rotateL Word64
b ([Int]
rots [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
8)) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
x
            t :: Int
t = (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
4
            (Word64
p, Word64
q) =
                if (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                    then (Word64
x Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ [Word64]
keys [Word64] -> Int -> Word64
forall a. HasCallStack => [a] -> Int -> a
!! (Int
t Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
3), Word64
y Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ [Word64]
keys [Word64] -> Int -> Word64
forall a. HasCallStack => [a] -> Int -> a
!! ((Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
3) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
t)
                    else (Word64
x, Word64
y)
         in Int -> Word64 -> Word64 -> (Word64, Word64)
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Word64
p Word64
q
chacha :: Word64 -> [Word64] -> [Word32]
chacha :: Word64 -> [Word64] -> [Word32]
chacha Word64
c [Word64]
ks = (Word32 -> Word32 -> Word32) -> [Word32] -> [Word32] -> [Word32]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
(+) [Word32]
initial (([Word32] -> [Word32]) -> [Word32] -> [[Word32]]
forall a. (a -> a) -> a -> [a]
iterate [Word32] -> [Word32]
forall {a}. (Num a, Bits a) => [a] -> [a]
double [Word32]
initial [[Word32]] -> Int -> [Word32]
forall a. HasCallStack => [a] -> Int -> a
!! Int
10)
  where
    initial :: [Word32]
initial =
        [Word32
0x61707865, Word32
0x3320646e, Word32
0x79622d32, Word32
0x6b206574]
            [Word32] -> [Word32] -> [Word32]
forall a. [a] -> [a] -> [a]
++ (Word64 -> [Word32]) -> [Word64] -> [Word32]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\Word64
k -> [Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
k, Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
k Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
32)]) [Word64]
ks
            [Word32] -> [Word32] -> [Word32]
forall a. [a] -> [a] -> [a]
++ [Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
c, Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
c Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
32), Word32
0, Word32
0]
    set :: Int -> a -> [a] -> [a]
set Int
j a
v [a]
xs = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
j [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
v] [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
xs
    quarter :: [a] -> (Int, Int, Int, Int) -> [a]
quarter [a]
xs (Int
a, Int
b, Int
d, Int
e) =
        let a1 :: a
a1 = [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
a a -> a -> a
forall a. Num a => a -> a -> a
+ [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
b
            e1 :: a
e1 = a -> Int -> a
forall a. Bits a => a -> Int -> a
rotateL ([a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
e a -> a -> a
forall a. Bits a => a -> a -> a
`xor` a
a1) Int
16
            d1 :: a
d1 = [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
d a -> a -> a
forall a. Num a => a -> a -> a
+ a
e1
            b1 :: a
b1 = a -> Int -> a
forall a. Bits a => a -> Int -> a
rotateL ([a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
b a -> a -> a
forall a. Bits a => a -> a -> a
`xor` a
d1) Int
12
            a2 :: a
a2 = a
a1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
b1
            e2 :: a
e2 = a -> Int -> a
forall a. Bits a => a -> Int -> a
rotateL (a
e1 a -> a -> a
forall a. Bits a => a -> a -> a
`xor` a
a2) Int
8
            d2 :: a
d2 = a
d1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
e2
            b2 :: a
b2 = a -> Int -> a
forall a. Bits a => a -> Int -> a
rotateL (a
b1 a -> a -> a
forall a. Bits a => a -> a -> a
`xor` a
d2) Int
7
         in Int -> a -> [a] -> [a]
forall {a}. Int -> a -> [a] -> [a]
set Int
e a
e2 (Int -> a -> [a] -> [a]
forall {a}. Int -> a -> [a] -> [a]
set Int
d a
d2 (Int -> a -> [a] -> [a]
forall {a}. Int -> a -> [a] -> [a]
set Int
b a
b2 (Int -> a -> [a] -> [a]
forall {a}. Int -> a -> [a] -> [a]
set Int
a a
a2 [a]
xs)))
    double :: [a] -> [a]
double [a]
xs =
        ([a] -> (Int, Int, Int, Int) -> [a])
-> [a] -> [(Int, Int, Int, Int)] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl'
            [a] -> (Int, Int, Int, Int) -> [a]
forall {a}. (Num a, Bits a) => [a] -> (Int, Int, Int, Int) -> [a]
quarter
            [a]
xs
            [ (Int
0, Int
4, Int
8, Int
12)
            , (Int
1, Int
5, Int
9, Int
13)
            , (Int
2, Int
6, Int
10, Int
14)
            , (Int
3, Int
7, Int
11, Int
15)
            , (Int
0, Int
5, Int
10, Int
15)
            , (Int
1, Int
6, Int
11, Int
12)
            , (Int
2, Int
7, Int
8, Int
13)
            , (Int
3, Int
4, Int
9, Int
14)
            ]

-- | Create two deterministic child streams by hashing two successive parent words. The original immutable parent is unchanged; see the module-level non-overlap caveat.
splitRNG :: RNG -> (RNG, RNG)
splitRNG :: RNG -> (RNG, RNG)
splitRNG r :: RNG
r@(RNG Generator
g [Word64]
_ [Word64]
_) =
    let (Word64
a, RNG
r1) = RNG -> (Word64, RNG)
nextWord64 RNG
r; (Word64
b, RNG
_) = RNG -> (Word64, RNG)
nextWord64 RNG
r1 in (Generator -> Word64 -> RNG
seedRNG Generator
g (Word64 -> Word64
mix Word64
a), Generator -> Word64 -> RNG
seedRNG Generator
g (Word64 -> Word64
mix (Word64
b Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
1)))

-- | Serialize full state, including pending block output, with the versioned hmath-rng-1 prefix.
serializeRNG :: RNG -> String
serializeRNG :: RNG -> String
serializeRNG (RNG Generator
g [Word64]
state [Word64]
buffered) = String
"hmath-rng-1 " String -> ShowS
forall a. [a] -> [a] -> [a]
++ RawRNG -> String
forall a. Show a => a -> String
show (Generator -> [Word64] -> [Word64] -> RawRNG
RawRNG Generator
g [Word64]
state [Word64]
buffered)

-- | Parse a versioned full state and validate generator lengths, forbidden zero states, and MT word/index ranges.
restoreRNG :: String -> Either MathError RNG
restoreRNG :: String -> Either MathError RNG
restoreRNG String
text = case Int -> String -> (String, String)
forall a. Int -> [a] -> ([a], [a])
splitAt Int
12 String
text of
    (String
"hmath-rng-1 ", String
rest) -> case String -> Maybe RawRNG
forall a. Read a => String -> Maybe a
readMaybe String
rest of
        Just (RawRNG Generator
g [Word64]
s [Word64]
ws) | Generator -> [Word64] -> Bool
forall {a}. Integral a => Generator -> [a] -> Bool
valid Generator
g [Word64]
s Bool -> Bool -> Bool
&& [Word64] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word64]
ws Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
8 -> RNG -> Either MathError RNG
forall a b. b -> Either a b
Right (Generator -> [Word64] -> [Word64] -> RNG
RNG Generator
g [Word64]
s [Word64]
ws)
        Maybe RawRNG
_ -> MathError -> Either MathError RNG
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"invalid RNG state")
    (String, String)
_ -> MathError -> Either MathError RNG
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"unsupported RNG state version")
  where
    valid :: Generator -> [a] -> Bool
valid Generator
PCG32 [a]
s = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 Bool -> Bool -> Bool
&& a -> Bool
forall a. Integral a => a -> Bool
odd ([a] -> a
forall a. HasCallStack => [a] -> a
last [a]
s)
    valid Generator
Xoshiro256StarStar [a]
s = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
4 Bool -> Bool -> Bool
&& (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0) [a]
s
    valid Generator
Xoroshiro128Plus [a]
s = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 Bool -> Bool -> Bool
&& (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0) [a]
s
    valid Generator
MT19937 [a]
s = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
625 Bool -> Bool -> Bool
&& [a] -> a
forall a. HasCallStack => [a] -> a
last [a]
s a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
624 Bool -> Bool -> Bool
&& (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0xffffffff) (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
624 [a]
s)
    valid Generator
Philox4x32 [a]
s = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2
    valid Generator
Threefry2x64 [a]
s = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2
    valid Generator
ChaCha20 [a]
s = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
5

-- | Uniform 53-bit real on [0,1), using the most significant random bits.
uniform :: RNG -> (Double, RNG)
uniform :: RNG -> (Double, RNG)
uniform RNG
r = let (Word64
w, RNG
s) = RNG -> (Word64, RNG)
nextWord64 RNG
r in (Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
w Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
11) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
9007199254740992, RNG
s)

-- | Uniform 52-bit midpoint real on (0,1), safe as input to logarithms.
uniformOpen :: RNG -> (Double, RNG)
uniformOpen :: RNG -> (Double, RNG)
uniformOpen RNG
r = let (Word64
w, RNG
s) = RNG -> (Word64, RNG)
nextWord64 RNG
r in ((Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
w Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
12) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
4503599627370496, RNG
s)

-- | Unbiased inclusive integer range using 64-bit rejection. Supports the complete machine Int range; rejects reversed bounds.
uniformInt :: Int -> Int -> RNG -> Either MathError (Int, RNG)
uniformInt :: Int -> Int -> RNG -> Either MathError (Int, RNG)
uniformInt Int
lo Int
hi RNG
r
    | Int
lo Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
hi = MathError -> Either MathError (Int, RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"uniformInt: lower > upper")
    | Bool
otherwise = Int -> RNG -> Either MathError (Int, RNG)
go Int
10000 RNG
r
  where
    width :: Integer
width = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
hi Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
lo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1
    bound :: Integer
bound = (Integer
2 Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
64 :: Int) :: Integer) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- ((Integer
2 Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
64 :: Int)) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
width)
    go :: Int -> RNG -> Either MathError (Int, RNG)
    go :: Int -> RNG -> Either MathError (Int, RNG)
go Int
0 RNG
_ = MathError -> Either MathError (Int, RNG)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"uniform integer rejection" Int
10000)
    go Int
n RNG
s =
        let (Word64
x, RNG
t) = RNG -> (Word64, RNG)
nextWord64 RNG
s
         in if Word64 -> Integer
forall a. Integral a => a -> Integer
toInteger Word64
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
bound
                then (Int, RNG) -> Either MathError (Int, RNG)
forall a b. b -> Either a b
Right (Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
lo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Word64 -> Integer
forall a. Integral a => a -> Integer
toInteger Word64
x Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
width), RNG
t)
                else Int -> RNG -> Either MathError (Int, RNG)
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) RNG
t

-- | Run a checked sampler sequentially, returning values and final state. Count must lie in [0,10000000].
sampleN :: Int -> (RNG -> Either MathError (a, RNG)) -> RNG -> Either MathError ([a], RNG)
sampleN :: forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
n RNG -> Either MathError (a, RNG)
f RNG
r
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10000000 = MathError -> Either MathError ([a], RNG)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sample count outside [0,10000000]")
    | Bool
otherwise = Int -> RNG -> [a] -> Either MathError ([a], RNG)
forall {t}.
(Eq t, Num t) =>
t -> RNG -> [a] -> Either MathError ([a], RNG)
go Int
n RNG
r []
  where
    go :: t -> RNG -> [a] -> Either MathError ([a], RNG)
go t
0 RNG
s [a]
a = ([a], RNG) -> Either MathError ([a], RNG)
forall a b. b -> Either a b
Right ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
a, RNG
s); go t
k RNG
s [a]
a = do (x, t) <- RNG -> Either MathError (a, RNG)
f RNG
s; go (k - 1) t (x : a)

finite :: Double -> Bool
finite :: Double -> Bool
finite Double
x = Bool -> Bool
not (Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
x)
positive :: Double -> Bool
positive :: Double -> Bool
positive Double
x = Double -> Bool
finite Double
x Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0

-- | Standard normal through the polar-coordinate Box--Muller transform, consuming two uniforms.
boxMuller :: RNG -> (Double, RNG)
boxMuller :: RNG -> (Double, RNG)
boxMuller RNG
r = let (Double
u, RNG
s) = RNG -> (Double, RNG)
uniformOpen RNG
r; (Double
v, RNG
t) = RNG -> (Double, RNG)
uniform RNG
s in (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
u) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
cos (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v), RNG
t)

-- | Standard normal through Marsaglia polar disk rejection, with 10000-proposal budget.
marsagliaNormal :: RNG -> Either MathError (Double, RNG)

-- | Standard normal through Marsaglia polar disk rejection, with 10000-proposal budget.
marsagliaNormal :: RNG -> Either MathError (Double, RNG)
marsagliaNormal = Int -> RNG -> Either MathError (Double, RNG)
forall {t}.
(Eq t, Num t) =>
t -> RNG -> Either MathError (Double, RNG)
go (Int
10000 :: Int)
  where
    go :: t -> RNG -> Either MathError (Double, RNG)
go t
0 RNG
_ = MathError -> Either MathError (Double, RNG)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Marsaglia polar" Int
10000)
    go t
n RNG
r =
        let (Double
u, RNG
s) = RNG -> (Double, RNG)
uniform RNG
r; (Double
v, RNG
t) = RNG -> (Double, RNG)
uniform RNG
s; x :: Double
x = Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1; y :: Double
y = Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1; q :: Double
q = Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y
         in if Double
q Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 Bool -> Bool -> Bool
&& Double
q Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1 then (Double, RNG) -> Either MathError (Double, RNG)
forall a b. b -> Either a b
Right (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* 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
q Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
q), RNG
t) else t -> RNG -> Either MathError (Double, RNG)
go (t
n t -> t -> t
forall a. Num a => a -> a -> a
- t
1) RNG
t

{- | Marsaglia--Tsang 128-layer Ziggurat normal generator. Horizontal
rectangles have equal area; the base includes an exponentially proposed tail.
Published normal cutoff/area constants: https://doi.org/10.18637/jss.v005.i08
Tables are generated by the area recurrence, not embedded implementation data.
-}
zigguratNormal :: RNG -> Either MathError (Double, RNG)
zigguratNormal :: RNG -> Either MathError (Double, RNG)
zigguratNormal = Int -> RNG -> Either MathError (Double, RNG)
forall {t}.
(Eq t, Num t) =>
t -> RNG -> Either MathError (Double, RNG)
draw (Int
100000 :: Int)
  where
    cutoff :: Double
cutoff = Double
3.442619855899 :: Double
    area :: Double
area = Double
0.00991256303526217 :: Double
    scale :: Double
scale = Double
2147483648 :: Double
    topDensity :: Double
topDensity = Double -> Double
forall a. Floating a => a -> a
exp (-Double
cutoff Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
cutoff Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2)
    width0 :: Double
width0 = Double
area Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
topDensity
    -- Descending widths x_127,...,x_1; x_0 is the full base width.
    descending :: [Double]
descending = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
127 ((Double -> Double) -> Double -> [Double]
forall a. (a -> a) -> a -> [a]
iterate (\Double
x -> 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
area Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
exp (-Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2)))) Double
cutoff)
    widths :: [Double]
widths = Double
width0 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double] -> [Double]
forall a. [a] -> [a]
reverse [Double]
descending
    heights :: [Double]
heights = Double
1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
x -> Double -> Double
forall a. Floating a => a -> a
exp (-Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2)) (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 [Double]
widths)
    thresholds :: [Integer]
thresholds =
        Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double
cutoff Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
width0 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
scale)
            Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: Integer
0
            Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: [Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor ([Double]
widths [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ [Double]
widths [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
scale) | Int
i <- [Int
2 .. Int
127]] ::
            [Integer]
    draw :: t -> RNG -> Either MathError (Double, RNG)
draw t
0 RNG
_ = MathError -> Either MathError (Double, RNG)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Ziggurat normal" Int
100000)
    draw t
budget RNG
state =
        let (Word64
word, RNG
next) = RNG -> (Word64, RNG)
nextWord64 RNG
state
            signed :: Integer
signed = Int32 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
word :: Int32) :: Integer
            layer :: Int
layer = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
word Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
127)
            x :: Double
x = Integer -> Double
forall a. Num a => Integer -> a
fromInteger Integer
signed Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double]
widths [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
layer Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
scale
         in if Integer -> Integer
forall a. Num a => a -> a
abs Integer
signed Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< [Integer]
thresholds [Integer] -> Int -> Integer
forall a. HasCallStack => [a] -> Int -> a
!! Int
layer
                then (Double, RNG) -> Either MathError (Double, RNG)
forall a b. b -> Either a b
Right (Double
x, RNG
next)
                else
                    if Int
layer Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                        then t -> Bool -> RNG -> Either MathError (Double, RNG)
forall {t}.
(Eq t, Num t) =>
t -> Bool -> RNG -> Either MathError (Double, RNG)
tailDraw (t
budget t -> t -> t
forall a. Num a => a -> a -> a
- t
1) (Integer
signed Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0) RNG
next
                        else
                            let (Double
u, RNG
lastState) = RNG -> (Double, RNG)
uniform RNG
next
                                ordinate :: Double
ordinate = [Double]
heights [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
layer Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
* ([Double]
heights [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
layer Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double]
heights [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
layer)
                             in if Double
ordinate Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double -> Double
forall a. Floating a => a -> a
exp (-Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) then (Double, RNG) -> Either MathError (Double, RNG)
forall a b. b -> Either a b
Right (Double
x, RNG
lastState) else t -> RNG -> Either MathError (Double, RNG)
draw (t
budget t -> t -> t
forall a. Num a => a -> a -> a
- t
1) RNG
lastState
    tailDraw :: t -> Bool -> RNG -> Either MathError (Double, RNG)
tailDraw t
0 Bool
_ RNG
_ = MathError -> Either MathError (Double, RNG)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Ziggurat exponential tail" Int
100000)
    tailDraw t
budget Bool
negative RNG
state =
        let (Double
u, RNG
next) = RNG -> (Double, RNG)
uniformOpen RNG
state
            (Double
v, RNG
lastState) = RNG -> (Double, RNG)
uniformOpen RNG
next
            x :: Double
x = -Double -> Double
forall a. Floating a => a -> a
log Double
u Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
cutoff
            y :: Double
y = -Double -> Double
forall a. Floating a => a -> a
log Double
v
         in if Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y 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 (Double, RNG) -> Either MathError (Double, RNG)
forall a b. b -> Either a b
Right ((if Bool
negative then -Double
1 else Double
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
cutoff Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x), RNG
lastState)
                else t -> Bool -> RNG -> Either MathError (Double, RNG)
tailDraw (t
budget t -> t -> t
forall a. Num a => a -> a -> a
- t
1) Bool
negative RNG
lastState

-- | Normal sample with finite mean and strictly positive standard deviation. Arithmetic overflow returns NonFiniteResult.
normal :: Double -> Double -> RNG -> Either MathError (Double, RNG)
normal :: Double -> Double -> RNG -> Either MathError (Double, RNG)
normal Double
mu Double
sd RNG
r
    | Bool -> Bool
not (Double -> Bool
finite Double
mu Bool -> Bool -> Bool
&& Double -> Bool
positive Double
sd) = MathError -> Either MathError (Double, RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"normal: finite mean and positive sd required")
    | Bool
otherwise =
        let (Double
z, RNG
s) = RNG -> (Double, RNG)
boxMuller RNG
r; x :: Double
x = Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
sd Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
z in if Double -> Bool
finite Double
x then (Double, RNG) -> Either MathError (Double, RNG)
forall a b. b -> Either a b
Right (Double
x, RNG
s) else MathError -> Either MathError (Double, RNG)
forall a b. a -> Either a b
Left MathError
NonFiniteResult

-- | Exponential sample parameterized by strictly positive rate, via an open-uniform logarithm.
exponential :: Double -> RNG -> Either MathError (Double, RNG)
exponential :: Double -> RNG -> Either MathError (Double, RNG)
exponential Double
rate RNG
r
    | Bool -> Bool
not (Double -> Bool
positive Double
rate) = MathError -> Either MathError (Double, RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"exponential: positive rate required")
    | Bool
otherwise =
        let (Double
u, RNG
s) = RNG -> (Double, RNG)
uniformOpen RNG
r; x :: Double
x = -Double -> Double
forall a. Floating a => a -> a
log Double
u Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
rate
         in if Double -> Bool
finite Double
x then (Double, RNG) -> Either MathError (Double, RNG)
forall a b. b -> Either a b
Right (Double
x, RNG
s) else MathError -> Either MathError (Double, RNG)
forall a b. a -> Either a b
Left MathError
NonFiniteResult

-- | Gamma shape/scale sample using Marsaglia--Tsang rejection and the shape-augmentation identity below shape one; 10000-proposal budget.
gamma :: Double -> Double -> RNG -> Either MathError (Double, RNG)
gamma :: Double -> Double -> RNG -> Either MathError (Double, RNG)
gamma Double
shape Double
scale RNG
r
    | Bool -> Bool
not (Double -> Bool
positive Double
shape Bool -> Bool -> Bool
&& Double -> Bool
positive Double
scale) =
        MathError -> Either MathError (Double, RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"gamma: positive shape and scale required")
    | Double
shape Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1 = do
        (x, s) <- Double -> Double -> RNG -> Either MathError (Double, RNG)
gamma (Double
shape Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1) Double
scale RNG
r
        let (u, t) = uniformOpen s
        pure (x * u ** (1 / shape), t)
    | Bool
otherwise = Int -> RNG -> Either MathError (Double, RNG)
forall {t}.
(Eq t, Num t) =>
t -> RNG -> Either MathError (Double, RNG)
go (Int
10000 :: Int) RNG
r
  where
    d :: Double
d = Double
shape 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)
    go :: t -> RNG -> Either MathError (Double, RNG)
go t
0 RNG
_ = MathError -> Either MathError (Double, RNG)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Marsaglia-Tsang gamma" Int
10000)
    go t
n RNG
s =
        let (Double
z, RNG
t) = RNG -> (Double, RNG)
boxMuller RNG
s; 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 -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
3 :: Int); (Double
u, RNG
w) = RNG -> (Double, RNG)
uniformOpen RNG
t
         in if Double
v Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 Bool -> Bool -> Bool
&& (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.0331 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
z Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
4 :: Int) 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
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
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))
                then if Double -> Bool
finite (Double
scale Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
d Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v) then (Double, RNG) -> Either MathError (Double, RNG)
forall a b. b -> Either a b
Right (Double
scale Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
d Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v, RNG
w) else MathError -> Either MathError (Double, RNG)
forall a b. a -> Either a b
Left MathError
NonFiniteResult
                else t -> RNG -> Either MathError (Double, RNG)
go (t
n t -> t -> t
forall a. Num a => a -> a -> a
- t
1) RNG
w

-- | Beta(a,b) as the ratio of independent gamma variates. Positive parameters required; unresolved underflow returns an error.
beta :: Double -> Double -> RNG -> Either MathError (Double, RNG)
beta :: Double -> Double -> RNG -> Either MathError (Double, RNG)
beta Double
a Double
b RNG
r = do
    (x, s) <- Double -> Double -> RNG -> Either MathError (Double, RNG)
gamma Double
a Double
1 RNG
r
    (y, t) <- gamma b 1 s
    let v = if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
y then Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
y Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
x) else let ratio :: Double
ratio = Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
y in Double
ratio Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
ratio)
    if finite v then pure (v, t) else Left NonFiniteResult

-- | Poisson count by exponential-arrival inversion. Rate is in [0,100000], work is capped at one million arrivals.
poisson :: Double -> RNG -> Either MathError (Int, RNG)
poisson :: Double -> RNG -> Either MathError (Int, RNG)
poisson Double
lambda RNG
r
    | Bool -> Bool
not (Double -> Bool
finite Double
lambda) Bool -> Bool -> Bool
|| Double
lambda Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 Bool -> Bool -> Bool
|| Double
lambda Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
100000 =
        MathError -> Either MathError (Int, RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Poisson rate must be in [0,100000]")
    | Double
lambda Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 = (Int, RNG) -> Either MathError (Int, RNG)
forall a b. b -> Either a b
Right (Int
0, RNG
r)
    | Bool
otherwise = Int -> Double -> RNG -> Either MathError (Int, RNG)
forall {t}.
(Ord t, Num t) =>
t -> Double -> RNG -> Either MathError (t, RNG)
go Int
0 Double
0 RNG
r
  where
    go :: t -> Double -> RNG -> Either MathError (t, RNG)
go !t
n !Double
t RNG
s
        | t
n t -> t -> Bool
forall a. Ord a => a -> a -> Bool
>= t
1000000 = MathError -> Either MathError (t, RNG)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Poisson arrivals" Int
1000000)
        | Bool
otherwise =
            let (Double
u, RNG
z) = RNG -> (Double, RNG)
uniformOpen RNG
s; q :: Double
q = Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
log Double
u in if Double
q Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
lambda then (t, RNG) -> Either MathError (t, RNG)
forall a b. b -> Either a b
Right (t
n, RNG
z) else t -> Double -> RNG -> Either MathError (t, RNG)
go (t
n t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) Double
q RNG
z

-- | Binomial count from n independent Bernoulli trials. n<=10000000 and p in [0,1].
binomial :: Int -> Double -> RNG -> Either MathError (Int, RNG)
binomial :: Int -> Double -> RNG -> Either MathError (Int, RNG)
binomial Int
n Double
p RNG
r
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10000000 Bool -> Bool -> Bool
|| Bool -> Bool
not (Double -> Bool
finite Double
p) Bool -> Bool -> Bool
|| 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 =
        MathError -> Either MathError (Int, RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"binomial: invalid n or probability")
    | Bool
otherwise = do (xs, s) <- Int
-> (RNG -> Either MathError (Double, RNG))
-> RNG
-> Either MathError ([Double], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
n ((Double, RNG) -> Either MathError (Double, RNG)
forall a b. b -> Either a b
Right ((Double, RNG) -> Either MathError (Double, RNG))
-> (RNG -> (Double, RNG)) -> RNG -> Either MathError (Double, RNG)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RNG -> (Double, RNG)
uniform) RNG
r; pure (length (filter (< p) xs), s)

-- | Draw an index in proportion to finite nonnegative weights; max scaling handles subnormal or overflowing totals and zero weights are never selected.
weighted :: [Double] -> RNG -> Either MathError (Int, RNG)
weighted :: [Double] -> RNG -> Either MathError (Int, RNG)
weighted [Double]
ws RNG
r
    | [Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
ws Bool -> Bool -> Bool
|| (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Double
x -> Bool -> Bool
not (Double -> Bool
finite Double
x) Bool -> Bool -> Bool
|| Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0) [Double]
ws Bool -> Bool -> Bool
|| Double
largest Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 =
        MathError -> Either MathError (Int, RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"weights must be finite, nonnegative with positive maximum")
    | Bool
otherwise = let (Double
u, RNG
s) = RNG -> (Double, RNG)
uniform RNG
r; target :: Double
target = Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
total in (Int, RNG) -> Either MathError (Int, RNG)
forall a b. b -> Either a b
Right (Double -> Int -> [Double] -> Int
forall {a} {t}. (Num a, Num t, Ord a, Ord t) => a -> t -> [a] -> t
select Double
target Int
0 [Double]
normalized, RNG
s)
  where
    largest :: Double
largest = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Double
0 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
ws)
    normalized :: [Double]
normalized = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
largest) [Double]
ws
    total :: Double
total = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
normalized
    -- The final positive entry is a safe rounding fallback. Zero entries are
    -- skipped even when a subnormal original total could have rounded u*sum up.
    select :: a -> t -> [a] -> t
select a
target t
i (a
w : [a]
rest)
        | a
w a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0 = a -> t -> [a] -> t
select a
target (t
i t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) [a]
rest
        | a
target a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
w = t
i
        | (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0) [a]
rest = a -> t -> [a] -> t
select (a
target a -> a -> a
forall a. Num a => a -> a -> a
- a
w) (t
i t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) [a]
rest
        | Bool
otherwise = t
i
    select a
_ t
i [] = t -> t -> t
forall a. Ord a => a -> a -> a
max t
0 (t
i t -> t -> t
forall a. Num a => a -> a -> a
- t
1)

-- | Opaque Walker alias table with normalized acceptance thresholds and alternative indices.
data AliasTable = AliasTable [Double] [Int] deriving (AliasTable -> AliasTable -> Bool
(AliasTable -> AliasTable -> Bool)
-> (AliasTable -> AliasTable -> Bool) -> Eq AliasTable
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: AliasTable -> AliasTable -> Bool
== :: AliasTable -> AliasTable -> Bool
$c/= :: AliasTable -> AliasTable -> Bool
/= :: AliasTable -> AliasTable -> Bool
Eq, Int -> AliasTable -> ShowS
[AliasTable] -> ShowS
AliasTable -> String
(Int -> AliasTable -> ShowS)
-> (AliasTable -> String)
-> ([AliasTable] -> ShowS)
-> Show AliasTable
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> AliasTable -> ShowS
showsPrec :: Int -> AliasTable -> ShowS
$cshow :: AliasTable -> String
show :: AliasTable -> String
$cshowList :: [AliasTable] -> ShowS
showList :: [AliasTable] -> ShowS
Show)

-- | Build a Walker alias table from finite nonnegative weights with positive maximum; max scaling avoids overflow.
aliasTable :: [Double] -> Either MathError AliasTable
aliasTable :: [Double] -> Either MathError AliasTable
aliasTable [Double]
ws = do
    _ <- [Double] -> RNG -> Either MathError (Int, RNG)
weighted [Double]
ws (Generator -> Word64 -> RNG
seedRNG Generator
PCG32 Word64
0)
    let n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ws
        largest = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Double]
ws
        scaled = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
largest) [Double]
ws
        total = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
scaled
        ps = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
x -> (Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
total) 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]
scaled
        small = [Int
i | (Int
i, Double
p) <- [Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Double]
ps, Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1]
        large = [Int
i | (Int
i, Double
p) <- [Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Double]
ps, Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
1]
    pure (build ps [0 .. n - 1] small large)
  where
    set :: Int -> a -> [a] -> [a]
set Int
i a
v [a]
xs = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
i [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
v] [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
xs
    build :: [Double] -> [Int] -> [Int] -> [Int] -> AliasTable
build [Double]
ps [Int]
as (Int
s : [Int]
ss) (Int
l : [Int]
ls) =
        let q :: Double
q = [Double]
ps [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
l Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
ps [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
s Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1; np :: [Double]
np = Int -> Double -> [Double] -> [Double]
forall {a}. Int -> a -> [a] -> [a]
set Int
l Double
q [Double]
ps; na :: [Int]
na = Int -> Int -> [Int] -> [Int]
forall {a}. Int -> a -> [a] -> [a]
set Int
s Int
l [Int]
as
         in if Double
q Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1 then [Double] -> [Int] -> [Int] -> [Int] -> AliasTable
build [Double]
np [Int]
na (Int
l Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
ss) [Int]
ls else [Double] -> [Int] -> [Int] -> [Int] -> AliasTable
build [Double]
np [Int]
na [Int]
ss (Int
l Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
ls)
    build [Double]
ps [Int]
as [Int]
ss [Int]
ls = [Double] -> [Int] -> AliasTable
AliasTable (([Double] -> Int -> [Double]) -> [Double] -> [Int] -> [Double]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (\[Double]
a Int
i -> Int -> Double -> [Double] -> [Double]
forall {a}. Int -> a -> [a] -> [a]
set Int
i Double
1 [Double]
a) [Double]
ps ([Int]
ss [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
ls)) [Int]
as

-- | Draw an index from an already validated alias table, using one uniform integer and one uniform real.
aliasSample :: AliasTable -> RNG -> Either MathError (Int, RNG)
aliasSample :: AliasTable -> RNG -> Either MathError (Int, RNG)
aliasSample (AliasTable [Double]
ps [Int]
as) RNG
r = do
    (i, s) <- Int -> Int -> RNG -> Either MathError (Int, RNG)
uniformInt Int
0 ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) RNG
r
    let (u, t) = uniform s
    pure (if u < ps !! i then i else as !! i, t)

-- | Fisher--Yates permutation of list positions. This immutable list reference implementation costs quadratic time.
shuffle :: [a] -> RNG -> Either MathError ([a], RNG)
shuffle :: forall a. [a] -> RNG -> Either MathError ([a], RNG)
shuffle [a]
xs RNG
r = Int -> [a] -> RNG -> Either MathError ([a], RNG)
forall {a}. Int -> [a] -> RNG -> Either MathError ([a], RNG)
go ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
xs RNG
r
  where
    set :: Int -> a -> [a] -> [a]
set Int
i a
v [a]
ys = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
i [a]
ys [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
v] [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
ys
    go :: Int -> [a] -> RNG -> Either MathError ([a], RNG)
go Int
i [a]
ys RNG
s
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = ([a], RNG) -> Either MathError ([a], RNG)
forall a b. b -> Either a b
Right ([a]
ys, RNG
s)
        | Bool
otherwise = do (j, t) <- Int -> Int -> RNG -> Either MathError (Int, RNG)
uniformInt Int
0 Int
i RNG
s; go (i - 1) (set j (ys !! i) (set i (ys !! j) ys)) t

-- | Draw count values independently from a nonempty finite population, preserving duplicates.
withReplacement :: Int -> [a] -> RNG -> Either MathError ([a], RNG)
withReplacement :: forall {a}. Int -> [a] -> RNG -> Either MathError ([a], RNG)
withReplacement Int
n [a]
xs RNG
r
    | [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = MathError -> Either MathError ([a], RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"empty sampling population")
    | Bool
otherwise = Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
n (\RNG
s -> do (i, t) <- Int -> Int -> RNG -> Either MathError (Int, RNG)
uniformInt Int
0 ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) RNG
s; pure (xs !! i, t)) RNG
r

-- | Uniform subset in shuffled order. Count cannot exceed population length; uniqueness is by position.
withoutReplacement :: Int -> [a] -> RNG -> Either MathError ([a], RNG)
withoutReplacement :: forall {a}. Int -> [a] -> RNG -> Either MathError ([a], RNG)
withoutReplacement Int
n [a]
xs RNG
r
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs = MathError -> Either MathError ([a], RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"sample size exceeds population")
    | Bool
otherwise = do (ys, s) <- [a] -> RNG -> Either MathError ([a], RNG)
forall a. [a] -> RNG -> Either MathError ([a], RNG)
shuffle [a]
xs RNG
r; pure (take n ys, s)

-- | Algorithm R reservoir over a finite stream, retaining k positions uniformly without replacement.
reservoir :: Int -> [a] -> RNG -> Either MathError ([a], RNG)
reservoir :: forall {a}. Int -> [a] -> RNG -> Either MathError ([a], RNG)
reservoir Int
k [a]
xs RNG
r
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs = MathError -> Either MathError ([a], RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"reservoir size outside population")
    | Bool
otherwise = Int -> [a] -> [a] -> RNG -> Either MathError ([a], RNG)
forall {a}. Int -> [a] -> [a] -> RNG -> Either MathError ([a], RNG)
go Int
k (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
k [a]
xs) (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
k [a]
xs) RNG
r
  where
    go :: Int -> [a] -> [a] -> RNG -> Either MathError ([a], RNG)
go Int
_ [a]
a [] RNG
s = ([a], RNG) -> Either MathError ([a], RNG)
forall a b. b -> Either a b
Right ([a]
a, RNG
s)
    go Int
i [a]
a (a
y : [a]
ys) RNG
s = do
        (j, t) <- Int -> Int -> RNG -> Either MathError (Int, RNG)
uniformInt Int
0 Int
i RNG
s; go (i + 1) (if j < k then take j a ++ [y] ++ drop (j + 1) a else a) ys t

-- | One uniform point per equal-width interval of [0,1), in increasing stratum order.
stratified :: Int -> RNG -> Either MathError ([Double], RNG)
stratified :: Int -> RNG -> Either MathError ([Double], RNG)
stratified Int
n RNG
r
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = MathError -> Either MathError ([Double], RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"positive stratum count required")
    | Bool
otherwise = do
        (us, s) <- Int
-> (RNG -> Either MathError (Double, RNG))
-> RNG
-> Either MathError ([Double], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
n ((Double, RNG) -> Either MathError (Double, RNG)
forall a b. b -> Either a b
Right ((Double, RNG) -> Either MathError (Double, RNG))
-> (RNG -> (Double, RNG)) -> RNG -> Either MathError (Double, RNG)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RNG -> (Double, RNG)
uniform) RNG
r
        pure (zipWith (\Int
i Double
u -> (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
u) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) [0 :: Int ..] us, s)

-- | Self-normalized importance estimate from values and nonnegative finite importance weights. This is a biased finite-sample estimator.
importance :: [Double] -> [Double] -> Either MathError Double
importance :: [Double] -> [Double] -> Either MathError Double
importance [Double]
values [Double]
weights
    | [Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
values
        Bool -> Bool -> Bool
|| [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
values Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
weights
        Bool -> Bool -> Bool
|| (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Bool -> Bool
not (Bool -> Bool) -> (Double -> Bool) -> Double -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Bool
finite) [Double]
values
        Bool -> Bool -> Bool
|| (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Double
w -> Bool -> Bool
not (Double -> Bool
finite Double
w) Bool -> Bool -> Bool
|| Double
w Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0) [Double]
weights
        Bool -> Bool -> Bool
|| Bool -> Bool
not (Double -> Bool
positive ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
weights)) =
        MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"importance: invalid weights or observations")
    | Bool
otherwise = Double -> Either MathError Double
forall a b. b -> Either a b
Right ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
values [Double]
weights) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
weights)

-- | Accept proposals with the supplied target/envelope probability in [0,1]; the first argument is a hard proposal budget.
rejection ::
    Int ->
    (Double -> Double) ->
    (RNG -> Either MathError (Double, RNG)) ->
    RNG ->
    Either MathError (Double, RNG)
rejection :: Int
-> (Double -> Double)
-> (RNG -> Either MathError (Double, RNG))
-> RNG
-> Either MathError (Double, RNG)
rejection Int
budget Double -> Double
ratio RNG -> Either MathError (Double, RNG)
proposal = Int -> RNG -> Either MathError (Double, RNG)
forall {t}.
(Ord t, Num t) =>
t -> RNG -> Either MathError (Double, RNG)
go Int
budget
  where
    go :: t -> RNG -> Either MathError (Double, RNG)
go t
n RNG
s
        | t
n t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
0 = MathError -> Either MathError (Double, RNG)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"rejection sampling" Int
budget)
        | Bool
otherwise = do
            (x, t) <- RNG -> Either MathError (Double, RNG)
proposal RNG
s
            let (u, z) = uniform t; q = Double -> Double
ratio Double
x
            if not (finite q) || q < 0 || q > 1
                then Left (DomainError "rejection ratio outside [0,1]")
                else if u < q then Right (x, z) else go (n - 1) z

-- | Apply a checked inverse CDF to an open uniform; propagates inverse-domain or nonconvergence errors.
inverseTransform :: (Double -> Either MathError Double) -> RNG -> Either MathError (Double, RNG)
inverseTransform :: (Double -> Either MathError Double)
-> RNG -> Either MathError (Double, RNG)
inverseTransform Double -> Either MathError Double
inv RNG
r = let (Double
u, RNG
s) = RNG -> (Double, RNG)
uniformOpen RNG
r in do x <- Double -> Either MathError Double
inv Double
u; pure (x, s)

{- | Symmetric random-walk Metropolis in log density. Returned chain includes
each post-proposal state. Non-finite log densities are rejected.
-}
metropolisHastings ::
    Int -> Double -> (Double -> Double) -> Double -> RNG -> Either MathError ([Double], RNG)
metropolisHastings :: Int
-> Double
-> (Double -> Double)
-> Double
-> RNG
-> Either MathError ([Double], RNG)
metropolisHastings Int
n Double
step Double -> Double
logp Double
initial RNG
r
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Bool -> Bool
not (Double -> Bool
positive Double
step) Bool -> Bool -> Bool
|| Bool -> Bool
not (Double -> Bool
finite Double
initial Bool -> Bool -> Bool
&& Double -> Bool
finite (Double -> Double
logp Double
initial)) =
        MathError -> Either MathError ([Double], RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Metropolis: invalid initial state or step")
    | Bool
otherwise = Int
-> Double -> RNG -> [Double] -> Either MathError ([Double], RNG)
forall {t}.
(Eq t, Num t) =>
t -> Double -> RNG -> [Double] -> Either MathError ([Double], RNG)
go Int
n Double
initial RNG
r []
  where
    go :: t -> Double -> RNG -> [Double] -> Either MathError ([Double], RNG)
go t
0 Double
_ RNG
s [Double]
a = ([Double], RNG) -> Either MathError ([Double], RNG)
forall a b. b -> Either a b
Right ([Double] -> [Double]
forall a. [a] -> [a]
reverse [Double]
a, RNG
s)
    go t
k Double
x RNG
s [Double]
a = do
        (y, t) <- Double -> Double -> RNG -> Either MathError (Double, RNG)
normal Double
x Double
step RNG
s
        let (u, z) = uniformOpen t; lp = Double -> Double
logp Double
y; q = if Double -> Bool
finite Double
lp 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
lp Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
logp Double
x then Double
y else Double
x
        go (k - 1) q z (q : a)

{- | A full Gibbs sweep uses caller-supplied exact conditional samplers in
coordinate order. Shape is checked after each conditional.
-}
gibbs ::
    Int ->
    [([Double] -> RNG -> Either MathError (Double, RNG))] ->
    [Double] ->
    RNG ->
    Either MathError ([[Double]], RNG)
gibbs :: Int
-> [[Double] -> RNG -> Either MathError (Double, RNG)]
-> [Double]
-> RNG
-> Either MathError ([[Double]], RNG)
gibbs Int
n [[Double] -> RNG -> Either MathError (Double, RNG)]
fs [Double]
initial RNG
r
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| [[Double] -> RNG -> Either MathError (Double, RNG)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double] -> RNG -> Either MathError (Double, RNG)]
fs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
initial Bool -> Bool -> Bool
|| (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Bool -> Bool
not (Bool -> Bool) -> (Double -> Bool) -> Double -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Bool
finite) [Double]
initial =
        MathError -> Either MathError ([[Double]], RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Gibbs: invalid shape")
    | Bool
otherwise = Int
-> [Double]
-> RNG
-> [[Double]]
-> Either MathError ([[Double]], RNG)
forall {t}.
(Eq t, Num t) =>
t
-> [Double]
-> RNG
-> [[Double]]
-> Either MathError ([[Double]], RNG)
go Int
n [Double]
initial RNG
r []
  where
    sweep :: [Double] -> RNG -> Either MathError ([Double], RNG)
sweep [Double]
x RNG
s =
        (Either MathError ([Double], RNG)
 -> (Int, [Double] -> RNG -> Either MathError (Double, RNG))
 -> Either MathError ([Double], RNG))
-> Either MathError ([Double], RNG)
-> [(Int, [Double] -> RNG -> Either MathError (Double, RNG))]
-> Either MathError ([Double], RNG)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl'
            ( \Either MathError ([Double], RNG)
acc (Int
i, [Double] -> RNG -> Either MathError (Double, RNG)
f) -> do
                (y, t) <- Either MathError ([Double], RNG)
acc
                (v, z) <- f y t
                if finite v then Right (take i y ++ [v] ++ drop (i + 1) y, z) else Left NonFiniteResult
            )
            (([Double], RNG) -> Either MathError ([Double], RNG)
forall a b. b -> Either a b
Right ([Double]
x, RNG
s))
            ([Int]
-> [[Double] -> RNG -> Either MathError (Double, RNG)]
-> [(Int, [Double] -> RNG -> Either MathError (Double, RNG))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [[Double] -> RNG -> Either MathError (Double, RNG)]
fs)
    go :: t
-> [Double]
-> RNG
-> [[Double]]
-> Either MathError ([[Double]], RNG)
go t
0 [Double]
_ RNG
s [[Double]]
a = ([[Double]], RNG) -> Either MathError ([[Double]], RNG)
forall a b. b -> Either a b
Right ([[Double]] -> [[Double]]
forall a. [a] -> [a]
reverse [[Double]]
a, RNG
s)
    go t
k [Double]
x RNG
s [[Double]]
a = do (y, t) <- [Double] -> RNG -> Either MathError ([Double], RNG)
sweep [Double]
x RNG
s; go (k - 1) y t (y : a)

-- | Identity-mass HMC with Gaussian momentum, leapfrog count and positive step size, log density and its gradient. Returns post-Metropolis states; finite-gradient shapes are checked.
hamiltonianMonteCarlo ::
    Int ->
    Int ->
    Double ->
    ([Double] -> Double) ->
    ([Double] -> [Double]) ->
    [Double] ->
    RNG ->
    Either MathError ([[Double]], RNG)
hamiltonianMonteCarlo :: Int
-> Int
-> Double
-> ([Double] -> Double)
-> ([Double] -> [Double])
-> [Double]
-> RNG
-> Either MathError ([[Double]], RNG)
hamiltonianMonteCarlo Int
n Int
leaps Double
eps [Double] -> Double
logp [Double] -> [Double]
grad [Double]
initial RNG
r
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        Bool -> Bool -> Bool
|| Int
leaps Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
        Bool -> Bool -> Bool
|| Bool -> Bool
not (Double -> Bool
positive Double
eps)
        Bool -> Bool -> Bool
|| [Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
initial
        Bool -> Bool -> Bool
|| (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Bool -> Bool
not (Bool -> Bool) -> (Double -> Bool) -> Double -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Bool
finite) [Double]
initial
        Bool -> Bool -> Bool
|| Bool -> Bool
not (Double -> Bool
finite ([Double] -> Double
logp [Double]
initial)) =
        MathError -> Either MathError ([[Double]], RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"HMC: invalid configuration")
    | Bool
otherwise = Int
-> [Double]
-> RNG
-> [[Double]]
-> Either MathError ([[Double]], RNG)
forall {t}.
(Eq t, Num t) =>
t
-> [Double]
-> RNG
-> [[Double]]
-> Either MathError ([[Double]], RNG)
go Int
n [Double]
initial RNG
r []
  where
    add :: [c] -> [c] -> [c]
add [c]
a [c]
b = (c -> c -> c) -> [c] -> [c] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith c -> c -> c
forall a. Num a => a -> a -> a
(+) [c]
a [c]
b
    scale :: b -> [b] -> [b]
scale b
a = (b -> b) -> [b] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (b
a b -> b -> b
forall a. Num a => a -> a -> a
*)
    force :: [Double] -> Either MathError [Double]
force [Double]
q =
        let g :: [Double]
g = [Double] -> [Double]
grad [Double]
q
         in if [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
g Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
q Bool -> Bool -> Bool
|| (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Bool -> Bool
not (Bool -> Bool) -> (Double -> Bool) -> Double -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Bool
finite) [Double]
g
                then MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"HMC gradient shape or value")
                else [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right [Double]
g
    integrate :: t -> [Double] -> [Double] -> Either MathError ([Double], [Double])
integrate t
0 [Double]
q [Double]
p = ([Double], [Double]) -> Either MathError ([Double], [Double])
forall a b. b -> Either a b
Right ([Double]
q, [Double]
p)
    integrate t
k [Double]
q [Double]
p = do
        g <- [Double] -> Either MathError [Double]
force [Double]
q
        let ph = [Double] -> [Double] -> [Double]
forall {c}. Num c => [c] -> [c] -> [c]
add [Double]
p (Double -> [Double] -> [Double]
forall {b}. Num b => b -> [b] -> [b]
scale (Double
eps Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) [Double]
g); qn = [Double] -> [Double] -> [Double]
forall {c}. Num c => [c] -> [c] -> [c]
add [Double]
q (Double -> [Double] -> [Double]
forall {b}. Num b => b -> [b] -> [b]
scale Double
eps [Double]
ph)
        gn <- force qn
        integrate (k - 1) qn (add ph (scale (eps / 2) gn))
    go :: t
-> [Double]
-> RNG
-> [[Double]]
-> Either MathError ([[Double]], RNG)
go t
0 [Double]
_ RNG
s [[Double]]
a = ([[Double]], RNG) -> Either MathError ([[Double]], RNG)
forall a b. b -> Either a b
Right ([[Double]] -> [[Double]]
forall a. [a] -> [a]
reverse [[Double]]
a, RNG
s)
    go t
k [Double]
q RNG
s [[Double]]
a = do
        (p, t) <- Int
-> (RNG -> Either MathError (Double, RNG))
-> RNG
-> Either MathError ([Double], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
q) (Double -> Double -> RNG -> Either MathError (Double, RNG)
normal Double
0 Double
1) RNG
s
        (qn, pn) <- integrate leaps q p
        let (u, z) = uniformOpen t
            ratio = [Double] -> Double
logp [Double]
qn Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> Double
logp [Double]
q Double -> Double -> Double
forall a. Num a => a -> a -> a
+ ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
x -> Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x) [Double]
p) Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
x -> Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x) [Double]
pn)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
            chosen = if Double -> Bool
finite Double
ratio 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
ratio then [Double]
qn else [Double]
q
        go (k - 1) chosen z (chosen : a)

-- | Uniform-[0,1] Monte Carlo integral and estimated standard error from at least two draws. This is a stochastic estimate, not a certified bound.
monteCarlo :: (Double -> Double) -> Int -> RNG -> Either MathError ((Double, Double), RNG)
monteCarlo :: (Double -> Double)
-> Int -> RNG -> Either MathError ((Double, Double), RNG)
monteCarlo Double -> Double
f Int
n RNG
r
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = MathError -> Either MathError ((Double, Double), RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Monte Carlo needs at least 2 observations")
    | Bool
otherwise = do
        (us, s) <- Int
-> (RNG -> Either MathError (Double, RNG))
-> RNG
-> Either MathError ([Double], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
n ((Double, RNG) -> Either MathError (Double, RNG)
forall a b. b -> Either a b
Right ((Double, RNG) -> Either MathError (Double, RNG))
-> (RNG -> (Double, RNG)) -> RNG -> Either MathError (Double, RNG)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RNG -> (Double, RNG)
uniform) RNG
r
        let xs = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
f [Double]
us
            m = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
            se = Double -> Double
forall a. Floating a => a -> a
sqrt ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
x -> (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) [Double]
xs) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)))
        if all finite xs then Right ((m, se), s) else Left NonFiniteResult

-- | Unshifted Halton point, dimensions 1..16 and zero-based nonnegative point index; index zero is the origin.
halton :: Int -> Int -> Either MathError [Double]
halton :: Int -> Int -> Either MathError [Double]
halton Int
dimension Int
index
    | Int
dimension Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
dimension Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
16 Bool -> Bool -> Bool
|| Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 =
        MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Halton: dimension 1..16 and nonnegative index required")
    | Bool
otherwise =
        [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right
            [ Double -> Double -> Double
forall a. Ord a => a -> a -> a
min (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2 Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (-Double
53)) (Int -> Int -> Double
forall {t} {a}. (Fractional t, Integral a) => a -> a -> t
radical Int
b Int
index)
            | Int
b <- Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
dimension [Int
2, Int
3, Int
5, Int
7, Int
11, Int
13, Int
17, Int
19, Int
23, Int
29, Int
31, Int
37, Int
41, Int
43, Int
47, Int
53]
            ]
  where
    radical :: a -> a -> t
radical a
b = t -> t -> a -> t
forall {t}. Fractional t => t -> t -> a -> t
go (t
1 t -> t -> t
forall a. Fractional a => a -> a -> a
/ a -> t
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
b) t
0
      where
        go :: t -> t -> a -> t
go t
_ t
a a
0 = t
a; go t
w t
a a
k = t -> t -> a -> t
go (t
w t -> t -> t
forall a. Fractional a => a -> a -> a
/ a -> t
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
b) (t
a t -> t -> t
forall a. Num a => a -> a -> a
+ t
w t -> t -> t
forall a. Num a => a -> a -> a
* a -> t
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a
k a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
b)) (a
k a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
b)

{- | Sobol digital net in dimensions 1..8, 32 direction bits. Joe--Kuo's first
primitive polynomials and initial direction numbers; index zero is the origin.
-}
sobol :: Int -> Int -> Either MathError [Double]
sobol :: Int -> Int -> Either MathError [Double]
sobol Int
dimension Int
index
    | Int
dimension Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
dimension Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
8 Bool -> Bool -> Bool
|| Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
index Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
2 Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
32 :: Int) =
        MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Sobol: dimension 1..8, 32-bit index required")
    | Bool
otherwise = [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right [Int -> Double
forall {a}. Fractional a => Int -> a
coordinate Int
d | Int
d <- [Int
0 .. Int
dimension Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
  where
    specs :: [(Int, Int, [Int])]
    specs :: [(Int, Int, [Int])]
specs =
        [ (Int
1, Int
0, [Int
1])
        , (Int
2, Int
1, [Int
1, Int
3])
        , (Int
3, Int
1, [Int
1, Int
3, Int
1])
        , (Int
3, Int
2, [Int
1, Int
1, Int
1])
        , (Int
4, Int
1, [Int
1, Int
1, Int
3, Int
3])
        , (Int
4, Int
4, [Int
1, Int
3, Int
5, Int
13])
        , (Int
5, Int
2, [Int
1, Int
1, Int
5, Int
5, Int
17])
        ]
    dirs :: Int -> [Word32]
    dirs :: Int -> [Word32]
dirs Int
0 = [Int -> Word32
forall a. Bits a => Int -> a
bit (Int
31 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j) | Int
j <- [Int
0 .. Int
31]]
    dirs Int
d =
        let (Int
s, Int
a, [Int]
ms) = [(Int, Int, [Int])]
specs [(Int, Int, [Int])] -> Int -> (Int, Int, [Int])
forall a. HasCallStack => [a] -> Int -> a
!! (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            start :: [Word32]
start = (Word32 -> Int -> Word32) -> [Word32] -> [Int] -> [Word32]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
shiftL ((Int -> Word32) -> [Int] -> [Word32]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Int]
ms) [Int
31, Int
30 ..]
            go :: [a] -> Int -> [a]
go [a]
vs Int
j
                | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
32 = [a]
vs
                | Bool
otherwise =
                    let v :: a
v = [a]
vs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
s) a -> a -> a
forall a. Bits a => a -> a -> a
`xor` ([a]
vs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
s) a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
s)
                        w :: a
w =
                            (a -> Int -> a) -> a -> [Int] -> a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl'
                                (\a
x Int
k -> if Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit (Int
a :: Int) (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k) then a
x a -> a -> a
forall a. Bits a => a -> a -> a
`xor` [a]
vs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k) else a
x)
                                a
v
                                [Int
1 .. Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                     in [a] -> Int -> [a]
go ([a]
vs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
w]) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
         in [Word32] -> Int -> [Word32]
forall {a}. Bits a => [a] -> Int -> [a]
go [Word32]
start Int
s
    gray :: Int
gray = Int
index Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` (Int
index Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
1)
    coordinate :: Int -> a
coordinate Int
d =
        Word32 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word32 -> Word32 -> Word32) -> Word32 -> [Word32] -> Word32
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
xor Word32
0 [Word32
v | (Int
j, Word32
v) <- [Int] -> [Word32] -> [(Int, Word32)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] (Int -> [Word32]
dirs Int
d), Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
gray Int
j])
            a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
4294967296

-- | Faure's Pascal digital transform in the least prime base >= dimension.
faure :: Int -> Int -> Either MathError [Double]
faure :: Int -> Int -> Either MathError [Double]
faure Int
d Int
ix
    | Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
16 Bool -> Bool -> Bool
|| Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 =
        MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Faure dimension 1..16 and nonnegative index required")
    | Bool
otherwise = [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
d (([Int] -> Double) -> [[Int]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Double
forall {a} {a}. (Floating a, Integral a, Ord a) => [a] -> a
value (([Int] -> [Int]) -> [Int] -> [[Int]]
forall a. (a -> a) -> a -> [a]
iterate [Int] -> [Int]
transform [Int]
digits)))
  where
    b :: Int
b = Int -> Int
forall {t}. Integral t => t -> t
choosePrime (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
2 Int
d)
    choosePrime :: t -> t
choosePrime t
p =
        if (t -> Bool) -> [t] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\t
q -> t
p t -> t -> t
forall a. Integral a => a -> a -> a
`mod` t
q t -> t -> Bool
forall a. Eq a => a -> a -> Bool
/= t
0) [t
2 .. Double -> t
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Double
forall a. Floating a => a -> a
sqrt (t -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral t
p :: Double))]
            then t
p
            else t -> t
choosePrime (t
p t -> t -> t
forall a. Num a => a -> a -> a
+ t
1)
    digits :: [Int]
digits = if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then [Int
0] else Int -> [Int]
digs Int
ix
    digs :: Int -> [Int]
digs Int
0 = []; digs Int
n = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
b Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int -> [Int]
digs (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
b)
    choose :: a -> a -> a
choose a
n a
k =
        Integer -> a
forall a. Num a => Integer -> a
fromInteger
            ( ([Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [a -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
k a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) .. a -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n] Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Integer
1 .. a -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
k])
                Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
b
            )
    transform :: [Int] -> [Int]
transform [Int]
xs = [[Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int -> Int -> Int
forall {a} {a}. (Integral a, Num a) => a -> a -> a
choose Int
j Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* ([Int]
xs [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
j) | Int
j <- [Int
i .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
b | Int
i <- [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    value :: [a] -> a
value [a]
xs =
        a -> a -> a
forall a. Ord a => a -> a -> a
min
            (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
2 a -> a -> a
forall a. Floating a => a -> a -> a
** (-a
53))
            ([a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
b a -> a -> a
forall a. Floating a => a -> a -> a
** Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) | (Int
j, a
x) <- [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [a]
xs])

-- | Random n-by-d Latin hypercube in [0,1), with a separate Fisher--Yates stratum permutation per axis; n>0 and 1<=d<=1024.
latinHypercube :: Int -> Int -> RNG -> Either MathError ([[Double]], RNG)
latinHypercube :: Int -> Int -> RNG -> Either MathError ([[Double]], RNG)
latinHypercube Int
n Int
d RNG
r
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
|| Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
|| Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1024 = MathError -> Either MathError ([[Double]], RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Latin hypercube: invalid shape")
    | Bool
otherwise = do
        (cols, s) <- Int
-> (RNG -> Either MathError ([Double], RNG))
-> RNG
-> Either MathError ([[Double]], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
d (\RNG
t -> do (xs, u) <- Int -> RNG -> Either MathError ([Double], RNG)
stratified Int
n RNG
t; shuffle xs u) RNG
r
        pure (transpose cols, s)

{- | Mean, population variance, and lag-one correlation; descriptive quality
diagnostics, not a cryptographic certification or a replacement for TestU01.
-}
randomQuality :: Int -> RNG -> Either MathError ((Double, Double, Double), RNG)
randomQuality :: Int -> RNG -> Either MathError ((Double, Double, Double), RNG)
randomQuality Int
n RNG
r
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
3 = MathError -> Either MathError ((Double, Double, Double), RNG)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"quality sample count >=3 required")
    | Bool
otherwise = do
        (xs, s) <- Int
-> (RNG -> Either MathError (Double, RNG))
-> RNG
-> Either MathError ([Double], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
sampleN Int
n ((Double, RNG) -> Either MathError (Double, RNG)
forall a b. b -> Either a b
Right ((Double, RNG) -> Either MathError (Double, RNG))
-> (RNG -> (Double, RNG)) -> RNG -> Either MathError (Double, RNG)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RNG -> (Double, RNG)
uniform) RNG
r
        let m = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
            v = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
x -> (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) [Double]
xs) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
            c = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
x Double
y -> (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m)) [Double]
xs (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 [Double]
xs)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
v
        pure ((m, v, c), s)

{- | Read eight bytes from the operating system entropy device, little endian.
IO exceptions are propagated; this deliberately has no time-based fallback.
-}
entropySeed :: IO Word64
entropySeed :: IO Word64
entropySeed =
    String -> IOMode -> (Handle -> IO Word64) -> IO Word64
forall r. String -> IOMode -> (Handle -> IO r) -> IO r
withBinaryFile
        String
"/dev/urandom"
        IOMode
ReadMode
        ( \Handle
h -> do
            cs <- (Int -> IO Char) -> [Int] -> IO String
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\Int
_ -> Handle -> IO Char
hGetChar Handle
h) [Int
0 .. Int
7 :: Int]
            pure
                (Foldable.foldl' (\Word64
a (Int
i, Char
c) -> Word64
a Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Char -> Int
forall a. Enum a => a -> Int
fromEnum Char
c) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i)) 0 (zip [0 ..] cs))
        )