Faster looping with itertools

Question

I have a function

def getSamples():
    p = lambda x : mlab.normpdf(x,3,2) + mlab.normpdf(x,-5,1)
    q = lambda x : mlab.normpdf(x,5,14)
    k=30
    goodSamples = []
    rightCount = 0
    totalCount = 0
    while(rightCount < 100000):
        z0 = np.random.normal(5, 14)
        u0 = np.random.uniform(0,k*q(z0))
        if(p(z0) > u0):
            goodSamples.append(z0)
            rightCount += 1
        totalCount += 1
    return np.array(goodSamples)

My implementation to generate 100000 samples is taking much long. How can I make it fast with itertools or something similar?

Answer 1

In general, performance problems are caused by doing things the "wrong way". Numpy can be very fast when used as it is designed to be used, that is by exploiting its vector processing where these vectorized operations are handed off to compiled code. Two bad practices that come from other programing languages/approaches are

1. Loops: Whenever you think you need a loop stop and think. Most of the time you do not and in fact do not even want one. It is much faster both to write and run code without loops.
2. Memory allocation: Whenever you know the size of an object, preallocate space for it. Growing memory, particularly in Python lists, is very slow compared to the alternatives.

In this case it is easy to get (approximately) two orders of magnitude speedup; the tradeoff is more memory usage.

Below is some representative code, it is not meant to be blindly used. I have not even verified it produces the correct results. It is more or less a direct translation of your routine. It appears you are drawing random numbers from a probability distribution using the rejection method. There may be more efficient algorithms to do this for your probability distribution.

def getSamples2() :
    p = lambda x : mlab.normpdf(x,3,2) + mlab.normpdf(x,-5,1)
    q = lambda x : mlab.normpdf(x,5,14)
    k=30
    N = 100000 # Total number of samples we want
    Ngood = 0 # Current number of good samples
    goodSamples = np.zeros(N) # Storage for the good samples
    while Ngood < N : # Unfortunately a loop, ....
       z0 = np.random.normal(5, 14, size=N)
       u0 = np.random.uniform(size=N)*k*q(z0)
       ind, = np.where(p(z0) > u0)
       n = min(len(ind), N-Ngood)
       goodSamples[Ngood:Ngood+n] = z0[ind[:n]]
       Ngood += n
    return goodSamples

This generates random numbers in chunks and saves the good ones. I have not tried to optimize the chunk size (here I just use N, the total number we want, in principle this could/should be different and could even be adjusted based on the number we have left to generate). This still uses a loop, unfortunately, but now this will be run "tens" of times instead of 100,000 times. This also uses the where function and array slicing; these are good general tools to be comfortable with.

In one test with %timeit on my machine I found

In [27]: %timeit getSamples() # Original routine
1 loops, best of 3: 49.3 s per loop
In [28]: %timeit getSamples2()
1 loops, best of 3: 505 ms per loop

Answer 2

I would say that the secret to making this code faster does not lie in changing the loop syntax. Here are a few points:

1. np.random.normal has an additional parameter size that lets you get many values at once. I would suggest using an array of say 1E09 elements and then checking your condition on that for how many are good. You can then estimate how likely that is.
2. To create your uniform samples, why not use sympy for symbolic evaluation of the pdf? (I don't know if this is faster but it could be since you already know the mean and variance.)
3. Again, for p could you use a symbolic function?

Answer 3

Here is kinda itertools "magic", but I'm not sure it can help. Probably it's much better for perfomance to prepare an numpy array (using zeros) and fill it without creating python auto-growing list. Here is both itertools and zero-preparations. (Excuse me in advance for untested code)

from itertools import count, ifilter, imap, takewhile
import operator


def getSamples():
    p = lambda x : mlab.normpdf(x, 3, 2) + mlab.normpdf(x, -5, 1)
    q = lambda x : mlab.normpdf(x, 5, 14)

    k = 30
    n = 100000
    samples_iter = imap(
        operator.itemgetter(1),
        takewhile(
            lambda i, s: i < n,
            enumerate(
                ifilter(lambda z: p(z) > np.random.uniform(0,k*q(z)),
                        (np.random.normal(5, 14) for _ in count()))
    )))

    goodSamples = numpy.zeros(n)
    # set values from iterator, probably there is a better way for that
    for i, sample in enumerate(samples_iter):
        goodSamples[i] = sample
    return goodSamples