## How I Got a 327x Speedup Of Some Python Code

I don’t usually post code stuff on this blog, but I had some fun working on this and I wanted to share!

My colleague, Abhi, is processing data from an instrument collecting actual data from the real world (you should know that this is something I have never done!) and is having some problems with how long his analysis is taking. In particular, it is taking him longer than a day to analyze a day’s worth of data, which is clearly an unacceptable rate of progress. One step in his analysis is to take the data from all ~30,000 frequency channels of the instrument, and calculate an average across a moving window 1,000 entries long. For example, if I had the list:

^{?}View Code TEXT

[0, 1, 2, 3, 4, 5, 6, 7, 8, 9] |

and I wanted to find the average over a window of length three, I’d get:

^{?}View Code TEXT

[1, 2, 3, 4, 5, 6, 7, 8] |

Notice that I started by averaging `[0, 1, 2]=>1`

, as this is the first way to come up with a window of three values. Similarly, the last entry is `[7, 8, 9]=>8`

. Here is a simplified version of how Ahbi was doing it (all of the code snippets shown here are also available here):

^{?}View Code PYTHON

def avg0(arr, l): # arr - the list of values # l - the length of the window to average over new = [] for i in range(arr.size - l + 1): s = 0 for j in range(i, l + i): s += j new.append(s / l) return new |

This function is correct, but when we time this simple function (using an IPython Magic) we get:

^{?}View Code PYTHON

# a is [0, 1, 2, ..., 29997, 29998, 29999] a = np.arange(30000) %timeit avg0(a, 1000) 1 loops, best of 3: 3.19 s per loop |

Which isn’t so horrible in the absolute sense – 3 seconds isn’t that long compared to how much time we all waste per day. However, I neglected to mention above that Abhi must do this averaging for nearly 9,000 chunks of data per collecting day. This means that this averaging step alone takes about 8 hours of processing time per day of collected data. This is clearly taking too long!

The next thing I tried is to take advantage of Numpy, which is a Python package that enables much faster numerical analysis than with vanilla Python. In this function I’m essentially doing the same thing as above, but using Numpy methods, including pre-allocating space, summing values using the accelerated Numpy `.sum()`

function, and using the `xrange()`

iterator, which is somewhat faster than plain `range()`

:

^{?}View Code PYTHON

def avg1(arr, l): new = np.empty(arr.size - l + 1, dtype=arr.dtype) for i in xrange(arr.size - l + 1): new[i] = arr[i:l+i].sum() / l return new |

This provides a healthy speed up of about a factor of eight:

^{?}View Code PYTHON

%timeit avg1(a, 1000) 1 loops, best of 3: 405 ms per loop |

But this still isn’t fast enough; we’ve gone from 8 hours to just under an hour. We’d like to run this analysis in at most a few minutes. Below is an improved method that takes advantage of a queue, which is a programming construct that allows you to efficiently keep track of values you’ve seen before.

^{?}View Code PYTHON

def avg2(arr, l): new = np.empty(arr.size - l + 1, dtype=arr.dtype) d = deque() s = 0 i = 0 for value in arr: d.append(value) s += value if len(d) == l: new[i] = s / l i += 1 elif len(d) > l: s -= d.popleft() new[i] = s / l i += 1 return new |

And here we can see that we’ve cut our time down by another factor of 4:

^{?}View Code PYTHON

%timeit avg2(a, 1000) 10 loops, best of 3: 114 ms per loop |

This means that from 8 hours we’re now down to roughly 13 minutes. Getting better, but still not great. What else can we try? I’ve been looking for an excuse to try out Numba, which is a tool that is supposed to help speed up numerical analysis in Python, so I decided to give it a shot. What makes Numba attractive is that with a single additional line, Numba can take a function and dramatically speed it up by seamlessly converting it into C and compiling it when needed. So let’s try this on the first averaging function:

^{?}View Code PYTHON

@jit(argtypes=[int32[:], int32], restype=int32[:]) def avg0_numba(arr, l): new = [] for i in range(arr.size - l + 1): s = 0 for j in range(i, l + i): s += j new.append(s / l) return np.array(new) |

In the line beginning with `@jit`

, all I have to do is describe the input and the output types, and it handles the rest. And here’s the result:

^{?}View Code PYTHON

%timeit avg0_numba(a, 1000) 10 loops, best of 3: 21.6 ms per loop |

What is incredible here is that not only is this roughly 5 times faster than the queue method above, it’s a ridiculous 147 times faster than the original method and only one line has been added. We’ve now reduced 8 hours to about 4 minutes. Not bad!

Let’s try this on the second averaging method, which if you recall, was substantially better than the original method:

^{?}View Code PYTHON

@jit(argtypes=[int32[:], int32], restype=int32[:]) def avg1_numba(arr, l): new = np.empty(arr.size - l + 1, dtype=arr.dtype) for i in xrange(arr.size - l + 1): new[i] = arr[i:l+i].sum() / l return new %timeit avg1_numba(a, 1000) 1 loops, best of 3: 688 ms per loop |

That’s interesting! For some reason I don’t understand, this is actually slower than the un-optimized version of `avg1`

. Let’s see if Numba can speed up the queue method:

^{?}View Code PYTHON

@jit(argtypes=[int32[:], int32], restype=int32[:]) def avg2_numba(arr, l): new = np.empty(arr.size - l + 1, dtype=arr.dtype) d = deque() s = 0 i = 0 for value in arr: d.append(value) s += value if len(d) == l: new[i] = s / l i += 1 elif len(d) > l: s -= d.popleft() new[i] = s / l i += 1 return new %timeit avg2_numba(a, 1000) 10 loops, best of 3: 77.5 ms per loop |

This is somewhat better than before, but still not as fast as `avg0_numba`

, which comes in at roughly 20ms. But what if I really try hard to optimize the queue method by using only Numpy arrays?

^{?}View Code PYTHON

@jit(argtypes=[int32[:], int32], restype=int32[:]) def avg2_numba2(arr, l): new = np.empty(arr.size - l + 1, dtype=arr.dtype) d = np.empty(l + 1, dtype=arr.dtype) s = 0 i = 0 left = 0 right = 0 full = False for j in xrange(arr.size): d[right] = arr[j] s += arr[j] right = (right + 1) % (l+1) if not full and right == l: new[i] = s / l i += 1 full = True elif full: s -= d[left] left = (left + 1) % (l+1) new[i] = s / l i += 1 return new %timeit avg2_numba2(a, 1000) 100 loops, best of 3: 9.77 ms per loop |

A ha! That’s even faster, and our 3.19s are now down to 9.77ms, an improvement of 327 times. The original 8 hours are now reduced to less than two minutes. I’ve put all my code into an IPython Notebook, which you can download and run for yourself. I hope you’ve enjoyed this as much as I have!

**Update:**

After chatting with a couple of my colleagues, this appears to be the fastest way to do this kind of operation:

^{?}View Code PYTHON

from scipy.signal import fftconvolve import numpy as np a = np.arange(30000) b = np.ones(1000) / 1000. %timeit fftconvolve(a, b, 'valid') 100 loops, best of 3: 6.25 ms per loop |