pois_rand optimization

Author: sivasathyaseeelan

ext/JumpProcessesKernelAbstractionsExt.jl

@@ -134,7 +134,7 @@ end
 
         # Poisson sampling
         @inbounds for k in 1:num_jumps
-            counts[k] = pois_rand(rate_cache[k])
+            counts[k] = pois_rand(PassthroughRNG(), rate_cache[k])
         end
 
         # Apply changes
@@ -214,12 +214,17 @@ end
 # GPU-compatible Poisson sampling 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 = rng isa PassthroughRNG ? randexp() : randexp(rng)
+    c = randexp(rng)
     while c < λ
         n += 1
-        c += rng isa PassthroughRNG ? randexp() : randexp(rng)
+        c += randexp(rng)
     end
     return n
 end
@@ -232,43 +237,50 @@ end
 #
 #   For μ sufficiently large, (i.e. >= 10.0)
 #
+ad_rand(λ) = ad_rand(Random.GLOBAL_RNG, λ)
 function ad_rand(rng::AbstractRNG, λ)
-    λ = Float64(λ)
     s = sqrt(λ)
-    d = 6.0 * λ^2
+    d = 6 * λ^2
     L = floor(Int, λ - 1.1484)
-
-    G = λ + s * (rng isa PassthroughRNG ? randn() : randn(rng))
+    # Step N
+    G = λ + s * randn(rng)
 
     if G >= 0
         K = floor(Int, G)
+        # Step I
         if K >= L
             return K
         end
 
-        U = rng isa PassthroughRNG ? rand() : rand(rng)
+        # 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
-        E = rng isa PassthroughRNG ? randexp() : randexp(rng)
-        U = 2 * (rng isa PassthroughRNG ? rand() : rand(rng)) - 1
-        T_val = 1.8 + copysign(E, U)
-        if T_val <= -0.6744
+        # 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_val)
+        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
@@ -277,30 +289,30 @@ end
 
 # Procedure F
 function procf(λ, K::Int, s::Float64)
-    INV_SQRT_2PI = 0.3989422804014327  # 1/sqrt(2π)
+    # can be pre-computed, but does not seem to affect performance
+    INV_SQRT_2PI = inv(sqrt(2pi))
     ω = INV_SQRT_2PI / s
-    b1 = 1 / (24 * λ)
-    b2 = 0.3 * b1^2
-    c3 = b1 * b2 / 7
+    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 = -λ
-        log_py = K * log(λ) - loggamma(K + 1)  # log(K!) via loggamma
+        log_py = K * log(λ) - loggamma(K + 1) # log(K!) via loggamma
         py = exp(log_py)
     else
-        δ = 1 / (12 * K)
+        δ = inv(12) / K
         δ -= 4.8 * δ^3
         V = (λ - K) / K
-        px = K * log1pmx(V) - δ
+        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
+    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
@@ -316,16 +328,19 @@ Generates Poisson(λ) distributed random numbers using a fast polyalgorithm.
 ## Examples
 
 ```julia
-# Simple Poisson random which works on GPU
+# Simple Poisson random
 pois_rand(λ)
 
-# Using RNG
+# Using another RNG
 using RandomNumbers
 rng = Xorshifts.Xoroshiro128Plus()
 pois_rand(rng, λ)
+
+# Simple Poisson random on GPU
+pois_rand(PassthroughRNG(), λ)
 ```
 """
-pois_rand(λ) = λ < 6 ? count_rand(PassthroughRNG(), λ) : ad_rand(PassthroughRNG(), λ)
+pois_rand(λ) = pois_rand(Random.GLOBAL_RNG, λ)
 pois_rand(rng::AbstractRNG, λ) = λ < 6 ? count_rand(rng, λ) : ad_rand(rng, λ)
 
 end