Kemp and Kemp (1991) Poisson random variates implementation

General Usage Performance
1

Currently, the sampler in Distributions.jl for poisson random variates is PoissonADSampler, which implements the algorithm from J.H. Ahrens, U. Dieter (1982).

I found the paper, “Poisson Random Variate Generation” from 1991 by Kemp and Kemp which does a survey of older algorithms and proposes another, implemented below, which they show to be over twice as fast as (Ahrens and Dieter (1982)) for certain parameters.

I considered only lambda>=10, which is where both samplers are accurate, neither PoissonADSampler nor this one pass the Distributions.jl tests for smaller parameters (although I assume that can be tweaked).

My implementation is about twice as slow as PoissonADSampler, even excluding startup costs. I assume it was a faster algorithm at that time because it uses only one uniform random variate, and nowadays perhaps our PRNGs are faster and we are better at sampling from other distributions.

I haven’t actually seen any other implementation of this on the internet, presumably because it is slow, so I figured I would post it here to see if anyone can make it go faster, or finds it useful.

struct PoissonKempSampler{T<:Real} <: Sampleable{Univariate,Discrete} 
    λ::T
#these types can be made more general, just easier to leave like this for now
    r::Int 
    α::Float64 
    g::Float64
    x_max::Int
    p_r_λ::Float64
    function PoissonKempSampler(λ::T) where T
        r = trunc(Int, λ+0.5)
        α = λ - r
        g = 1/sqrt(2*π*r)
        x_max = trunc(Int,2λ + 30) #from paper
        p_r_λ =_p_r(g,r,α)
        new{T}(λ,r,α,g,x_max,p_r_λ)
    end
end

function _p_r(g,r,α)
    return g*(1 - (1/(12r + 0.5 + 293/(720r)))) * ((r + 2α/3 - α^2/4 - α^2/(18r))/
                                                    (r + 2α/3 + α^2/4 -  α^2/(18r)))
end
function _f_r_r(sampler)
    g = sampler.g
    r = sampler.r
    return 0.5 + 2*g*(1/3)*
        (1-(23/15)/
            (12r + (15/14) + ((30557/4508)/
                12r + (138134432/105880005))))
end

function _f_r_λ(sampler)
    r = sampler.r
    α = sampler.α
    p_r_λ = sampler.p_r_λ
    return _f_r_r(sampler) - α * p_r_λ * ((r + α/2 - α^2/60 - α^2/(20r))/
                                                (r + α/2 + α^2/20 -  α^2/(20r)))
end

function rand(rng::AbstractRNG, s::PoissonKempSampler)
    λ = s.λ
    r = s.r
    α = s.α
    g = s.g
    x_max = s.x_max
    p_r_λ = s.p_r_λ 
    u = rand(rng)
    if u <= 0.5
        return downward_search(u,p_r_λ,r,λ)
    elseif u >= 0.5 + 7g/6  
        return upward_search(u,p_r_λ,r,λ,x_max)
    else 
        f_r_λ = _f_r_λ(s)
        if u > f_r_λ
            return upward_search(u,p_r_λ,r,λ,x_max)
        else
            return downward_search(u,p_r_λ,r,λ)
        end
    end
end

@inline function downward_search(u,p_r_λ,r,λ)
    u<p_r_λ && return r
    p = p_r_λ
    for i = 0:r-1
        u -= p
        p *= (r - i)/λ
        u < p && return r - i - 1 
    end
    return 0
end

@inline function upward_search(u,p_r_λ,r,λ,x_max)
    u = 1- u
    p = p_r_λ
    for i = r+1:x_max
        p *= λ/i
        u < p && return i 
        u -= p
    end
end
3 Likes