replaced with pois_rand with PassthroughRNG

Author: sivasathyaseeelan

Project.toml

@@ -12,15 +12,13 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 PoissonRandom = "e409e4f3-bfea-5376-8464-e040bb5c01ab"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
-SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"

ext/JumpProcessesKernelAbstractionsExt.jl

@@ -3,9 +3,7 @@ module JumpProcessesKernelAbstractionsExt
 using JumpProcesses, SciMLBase
 using KernelAbstractions, Adapt
 using StaticArrays
-using Random
-using LogExpFunctions: log1pmx
-using SpecialFunctions: loggamma
+using PoissonRandom, Random
 
 function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem,
         alg::SimpleTauLeaping,
@@ -134,7 +132,7 @@ end
 
         # Poisson sampling
         @inbounds for k in 1:num_jumps
-            counts[k] = pois_rand(PassthroughRNG(), rate_cache[k])
+            counts[k] = pois_rand(PoissonRandom.PassthroughRNG(), rate_cache[k])
         end
 
         # Apply changes
@@ -211,138 +209,4 @@ function vectorized_solve(probs, prob::JumpProblem, alg::SimpleTauLeaping;
     return ts, us
 end
 
-export pois_rand, PassthroughRNG
-
-# GPU-compatible Poisson sampling via PassthroughRNG
-struct PassthroughRNG <: AbstractRNG end
-
-rand(rng::PassthroughRNG) = Random.rand()
-randexp(rng::PassthroughRNG) = Random.randexp()
-randn(rng::PassthroughRNG) = Random.randn()
-
-count_rand(λ) = count_rand(Random.GLOBAL_RNG, λ)
-function count_rand(rng::AbstractRNG, λ)
-    n = 0
-    c = randexp(rng)
-    while c < λ
-        n += 1
-        c += randexp(rng)
-    end
-    return n
-end
-
-# Algorithm from:
-#
-#   J.H. Ahrens, U. Dieter (1982)
-#   "Computer Generation of Poisson Deviates from Modified Normal Distributions"
-#   ACM Transactions on Mathematical Software, 8(2):163-179
-#
-#   For μ sufficiently large, (i.e. >= 10.0)
-#
-ad_rand(λ) = ad_rand(Random.GLOBAL_RNG, λ)
-function ad_rand(rng::AbstractRNG, λ)
-    s = sqrt(λ)
-    d = 6 * λ^2
-    L = floor(Int, λ - 1.1484)
-    # Step N
-    G = λ + s * randn(rng)
-
-    if G >= 0
-        K = floor(Int, G)
-        # Step I
-        if K >= L
-            return K
-        end
-
-        # Step S
-        U = rand(rng)
-        if d * U >= (λ - K)^3
-            return K
-        end
-
-        # Step P
-        px, py, fx, fy = procf(λ, K, s)
-
-        # Step Q
-        if fy * (1 - U) <= py * exp(px - fx)
-            return K
-        end
-    end
-
-    while true
-        # Step E
-        E = randexp(rng)
-        U = 2 * rand(rng) - 1
-        T = 1.8 + copysign(E, U)
-        if T <= -0.6744
-            continue
-        end
-
-        K = floor(Int, λ + s * T)
-        px, py, fx, fy = procf(λ, K, s)
-        c = 0.1069 / λ
-
-        # Step H
-        @fastmath if c * abs(U) <= py * exp(px + E) - fy * exp(fx + E)
-            return K
-        end
-    end
-end
-
-# Procedure F
-function procf(λ, K::Int, s::Float64)
-    # can be pre-computed, but does not seem to affect performance
-    INV_SQRT_2PI = inv(sqrt(2pi))
-    ω = INV_SQRT_2PI / s
-    b1 = inv(24) / λ
-    b2 = 0.3 * b1 * b1
-    c3 = inv(7) * b1 * b2
-    c2 = b2 - 15 * c3
-    c1 = b1 - 6 * b2 + 45 * c3
-    c0 = 1 - b1 + 3 * b2 - 15 * c3
-
-    if K < 10
-        px = -float(λ)
-        log_py = K * log(λ) - loggamma(K + 1) # log(K!) via loggamma
-        py = exp(log_py)
-    else
-        δ = inv(12) / K
-        δ -= 4.8 * δ^3
-        V = (λ - K) / K
-        px = K * log1pmx(V) - δ # avoids need for table
-        py = INV_SQRT_2PI / sqrt(K)
-    end
-    X = (K - λ + 0.5) / s
-    X2 = X^2
-    fx = X2 / -2 # missing negation in pseudo-algorithm, but appears in fortran code.
-    fy = ω * (((c3 * X2 + c2) * X2 + c1) * X2 + c0)
-    return px, py, fx, fy
-end
-
-"""
-```julia
-pois_rand(λ)
-pois_rand(rng::AbstractRNG, λ)
-```
-
-Generates Poisson(λ) distributed random numbers using a fast polyalgorithm.
-
-## Examples
-
-```julia
-# Simple Poisson random
-pois_rand(λ)
-
-# Using another RNG
-using RandomNumbers
-rng = Xorshifts.Xoroshiro128Plus()
-pois_rand(rng, λ)
-
-# Simple Poisson random on GPU
-pois_rand(PoissonRandom.PassthroughRNG(), λ)
-```
-"""
-pois_rand(λ) = pois_rand(Random.GLOBAL_RNG, λ)
-pois_rand(rng::AbstractRNG, λ) = λ < 6 ? count_rand(rng, λ) : ad_rand(rng, λ)
-
 end

test/gpu/regular_jumps.jl

@@ -4,7 +4,7 @@ using KernelAbstractions, Adapt, CUDA
 using StableRNGs
 rng = StableRNG(12345)
 
-Nsims = 100_000
+Nsims = 100
 
 # SIR model with influx
 let