The Fastest Poisson(1) in the West

Let's start off by scoping out the competition. I'll use julia just because I don't stand a chance with R since I'm actually competing with C in both cases. Interestingly, I'm competing against the same C code since the current version of julia's Distributions package uses RMath library. Let's see how it does.

using Distributions
using BenchmarkTools

function baseline_version()
    return rand(Poisson(1))
end

@benchmark baseline_version()

Julia has a nice BenchmarkTools package, so I'll just paste the results verbatim. We're mostly interested in the last four rows with the minimum, median, mean, and maximum times. Obviously lower is better.

It is important to note the units for the timings: we're talking nano-seconds. I typically don't bother with optimizing something like this since I'll have to run it 6*10¹⁰ times to make up a minute of programming time. But this is for fun and glory, so I'll do it now.

Finally, as a side note, we can do slightly better by removing the initialization of the Poisson distribution.

const dist = Poisson(1)
function baseline_no_initialization()
    return rand(dist)
end

@benchmark baseline_no_initialization()

When something like variable initialization starts making a difference you're far too deep into the weeds. But let's gather up our machetes and press forward.

So let's try to do this ourselves. We'll start off with the obvious approach: we draw a uniform random variable and invert the cdf.

function inverse_cdf()
    return quantile(dist, rand())
end

@benchmark inverse_cdf()

This is obviously slower; but it'll form a good building block. Let's write our own version.

function loop_version()
    ii = 0
    u = rand()
    while u < ccdf(dist, ii)
        ii += 1
    end
    return ii
end

@benchmark loop_version()

This is smarter than it looks: this is actually performing almost binary search. Straightforward binary search works by taking a sorted list and repeatedly testing the middle element. If the target value is equal to the middle element you return; if it's less than you apply binary search to the left side of the array, if it's greater than you apply it to the right. The key insight is that we can split in the array in half and find the value in log(n) time rather than linear time.

Obviously we have a infinite list of values so something is going to have to be modified (unless you have a good way of determining finding element \(\infty/2\)). What we do instead is that we test the median element first and then recurse on that.

Because the tail probabilities decay exponentially, it turns out the binary search progression is exactly the same as linear search except for the beginning, when you should test 1 instead of 0 first. However, when I coded that version up it was slower, so we'll just go in the regular linear order.

Needless to say, it doesn't matter what order we go in since we're not even in the same ballpark as the baseline. Of course, there's far too much post left for me to have left it at that.

Well, it's pretty obvious what's going wrong in the previous example: we're calculating the ccdf function far too often. It doesn't change, so we can move this outside the loop and put it in a lookup table. We have a nice mechanism for doing this that doesn't involve an ugly global constant: generated functions.

This also leads to the question of how far we should go; I say 100 is enough. Technically I'm drawing from a truncated Poisson distribution, but the events this far out have three digits in the exponent so I'm pretty safe.

@generated function loop_cached_version()
    N = 100
    ccdfs = [[ccdf(Poisson(1), ii) for ii in 0:N]; 1.0]

    eval = quote
        u = rand()
        ii = 1
        while u < $ccdfs[ii]
            ii += 1
        end
        return ii - 1
    end

    return eval
end

@benchmark loop_cached_version()

I'm pretty pleased with that benchmark.

Another approach is to use the Alias method. It turns the linear time loop into a constant time array lookup. Here's a quick implementation based on skimming the Wikipedia article.

@generated function alias()
    ## Hacky table generation; you could optimize this better
    N = 100

    U = [N * pdf(Poisson(1), ii) for ii = 0:(N-1)]
    K = [-1 for ii = 0:(N-1)]

    for ii = 1:(N-1)
        ii = indmax(U)
        jj = findfirst((U .< 1.0) & (K .== -1))
        K[jj] = ii - 1
        U[ii] = U[ii] - (1 - U[jj])
    end

    expr = quote
        u = rand()
        ii = floor(Int64, u * $N) + 1
        yy = $N*u + 1 - ii
        return yy < $U[ii] ? ii : $K[ii]
    end

    return expr
end

@benchmark alias()

I'm actually fairly surprised that this is so slow. I imagine it's the floor function which requires a type conversion (which explains the memory allocation). It could probably be more competitive if you just worked with the bits yourself. So you'd make a table with 128 entries and index with the first byte and compare with the rest.

Having mentioned working with bits directly brings up another topic. One problem with all of these method is that it's rather wasteful with its bits. In fact the expected number of bits that we need is just the entropy: 1.882. So generating 64 bit numbers is doing a lot more work than is necessary.

To really take advantage of the bits you might try something like a cached implementation. Sketched out it would look something like this:

## This can be tuned for optimal performance
_optimal_buffer_length = ??

type CachedPoisson
    buffer::Vector{Int}
    counter::Int
end

function CachedPoisson()
    buffer = zeros(Int, _optimal_buffer_length)
    counter = _optimal_buffer_length

    return CachedPoisson(buffer, counter)
end

function rand(::CachedPoisson)
    dist.counter += 1
    if dist.counter == _optimal_buffer_length + 1
        u = rand()
        ## Code goes here
        dist.counter = 1
    end
    return dist.buffer[dist.counter]
end

This wrings the most entropy out of the RNG as it can.

Doing it this way should be really fast; when the buffer is non-empty you're just doing a buffer lookup and when the buffer is empty you're doing a single random number generation followed by the usual costs of your usual method.

So it seems like you can do a great deal better than the regular Poisson implementation, at least if you can specialize.

So which of these methods did I end up using? The base implementation of course!