Faster convolution of probability density functions in Python

Question

Suppose the convolution of a general number of discrete probability density functions needs to be calculated. For the example below there are four distributions which take on values 0,1,2 with the specified probabilities:

import numpy as np
pdfs = np.array([[0.6,0.3,0.1],[0.5,0.4,0.1],[0.3,0.7,0.0],[1.0,0.0,0.0]])

The convolution can be found like this:

pdf = pdfs[0]        
for i in range(1,pdfs.shape[0]):
    pdf = np.convolve(pdfs[i], pdf)

The probabilities of seeing 0,1,...,8 are then given by

array([ 0.09 ,  0.327,  0.342,  0.182,  0.052,  0.007,  0.   ,  0.   ,  0.   ])

This part is the bottleneck in my code and it seems there must be something available to vectorize this operation. Does anyone have a suggestion for making it faster?

Alternatively, a solution where you could use

pdf1 = np.array([[0.6,0.3,0.1],[0.5,0.4,0.1]])
pdf2 = np.array([[0.3,0.7,0.0],[1.0,0.0,0.0]])
convolve(pd1,pd2) 

and get the pairwise convolutions

 array([[ 0.18,  0.51,  0.24,  0.07,  0.  ], 
        [ 0.5,  0.4,  0.1,  0. ,  0. ]])

would also help tremendously.

Answer 1

You can compute the convolution of all your PDFs efficiently using fast fourier transforms (FFTs): the key fact is that the FFT of the convolution is the product of the FFTs of the individual probability density functions. So transform each PDF, multiply the transformed PDFs together, and then perform the inverse transform. You'll need to pad each input PDF with zeros to the appropriate length to avoid effects from wraparound.

This should be reasonably efficient: if you have m PDFs, each containing n entries, then the time to compute the convolution using this method should grow as (m^2)n log(mn). The time is dominated by the FFTs, and we're effectively computing m + 1 independent FFTs (m forward transforms and one inverse transform), each of an array of length no greater than mn. But as always, if you want real timings you should profile.

Here's some code:

import numpy.fft

def convolve_many(arrays):
    """
    Convolve a list of 1d float arrays together, using FFTs.
    The arrays need not have the same length, but each array should
    have length at least 1.

    """
    result_length = 1 + sum((len(array) - 1) for array in arrays)

    # Copy each array into a 2d array of the appropriate shape.
    rows = numpy.zeros((len(arrays), result_length))
    for i, array in enumerate(arrays):
        rows[i, :len(array)] = array

    # Transform, take the product, and do the inverse transform
    # to get the convolution.
    fft_of_rows = numpy.fft.fft(rows)
    fft_of_convolution = fft_of_rows.prod(axis=0)
    convolution = numpy.fft.ifft(fft_of_convolution)

    # Assuming real inputs, the imaginary part of the output can
    # be ignored.
    return convolution.real

Applying this to your example, here's what I get:

>>> convolve_many([[0.6, 0.3, 0.1], [0.5, 0.4, 0.1], [0.3, 0.7], [1.0]])
array([ 0.09 ,  0.327,  0.342,  0.182,  0.052,  0.007])

That's the basic idea. If you want to tweak this, you might also look at numpy.fft.rfft (and its inverse, numpy.fft.irfft), which take advantage of the fact that the input is real to produce more compact transformed arrays. You might also be able to gain some speed by padding the rows array with zeros so that the total number of columns is optimal for performing an FFT. The definition of "optimal" here would depend on the FFT implementation, but powers of two would be good targets, for example. Finally, there are some obvious simplifications that can be made when creating rows if all the input arrays have the same length. But I'll leave these potential enhancements to you.