Weighted Random Choice

Question

I have a Cython function that takes a list of weights/probabilities (double) and returns a random index into the list.

For example:

choose_one(np.array([0.1, 0.4, 0.2, 0.3]), 4)

Returns 0 with probability 0.1, 1 with probability 0.4 etc.

The Cython code is:

cdef extern from "stdlib.h":
    double drand48()
    void srand48(long int seedval)

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
def choose_one(np.ndarray[double, ndim=1] weights, Py_ssize_t length):
    cdef Py_ssize_t idx, i
    cdef double cs
    cdef double random
    random = drand48()
    cs = 0.0
    i = 0
    while cs < random and i < length:
        cs += weights[i]
        i += 1
    return i - 1

I profiled some complex simulation code and found that about 50% of the execution time is spent in this function which is called approximately 3.5 million times. Because it is such a time hog, and such an important utility, I'd like to speed it up if possible.

One of the problems is that the input arrays are often very large with thousands of entries.

I tried various methods from Numpy in Cython and none of them come close to as fast as what I have written in Cython. I am a Cython newbie and am curious if there is something else I need to do to speed this up.

Answer 1

The performance can be improved a bit more by declaring the function using cpdef (or cdef) instead of def (when the calling code is using Cython types too, check out http://docs.cython.org/src/userguide/language_basics.html#python-functions-vs-c-functions for more info).

Another idea that might help is to check whether the random number is greater than 0.5 and, if it is, count backwards from the end of the array rather than the front (assuming that the values in the array always add up to one).

Also, if this function takes up half of the program's running time then it seems likely that the program could be restructured to use a different algorithm which calls this function less often, but this would probably require much more effort than the other suggestions.

Answer 2

If you understand enough C I'd suggest to look at the generated C code to judge whether it's not "optimal" yet. That said it looks like you've done everything for this single function; apart from switching to float instead of double for comparison I don't see anything in particular to change.

With that many invocations though the bottleneck will probably be the switch between Python and C, considering that the function itself does very little you'll not get very far by speeding up the function, instead you'll have to either do more on the C side, i.e. generate multiple results and pass them back (in which you should compare with numpy.random.choice I guess) or pass in more data; besides that you still have the option to parallelise this operation.