refactor

Author: sivasathyaseeelan

src/PoissonRandom.jl

@@ -4,17 +4,22 @@ using Random
 using LogExpFunctions: log1pmx
 using SpecialFunctions: loggamma
 
-export pois_rand
+export pois_rand, PassthroughRNG
 
 # 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
@@ -27,43 +32,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
@@ -89,13 +101,12 @@ function procf(λ, K::Int, s::Float64)
         δ = 1 / (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
@@ -120,7 +131,7 @@ rng = Xorshifts.Xoroshiro128Plus()
 pois_rand(rng, λ)
 ```
 """
-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 # module