Numpy performance games

I was reading a blog post about Numba vs Cython, which reported a huge speedup in using Numba with a simply Python code for calculating the pairwise distance between a number of points. The code they posted was as follows:

import numpy as np

def pairwise_python(X, D):
    M = X.shape[0]
    N = X.shape[1]
    for i in xrange(M):
        for j in xrange(M):
            d = 0.0
            for k in xrange(N):
                tmp = X[i, k] - X[j, k]
                d += tmp * tmp
            D[i, j] = np.sqrt(d)

They tested it with the following code:

N = 1000
X = np.random.random((N, 3))
D = np.empty((N, N))

%timeit pairwise_python(X, D)

 1 loops, best of 3: 8.3 s per loop

I thought I could probably do a bit better by making more use of the built-in numpy functions, so I rewrote it like thise:

def pairwise_numpy(X, D):
    M, N = X.shape
    for i in xrange(M):
        for j in xrange(M):
            D[i, j] = np.linalg.norm(X[i, :] - X[j, :])

%timeit pairwise_numpy(X, D)

1 loops, best of 3: 13.8 s per loop

Imagine my surprise. My “optimised” version takes almost twice as long! Shocked, I thought I shold try to go all the way

def pairwise_numpy_fast(X, D):
    for i, row in enumerate(X):
        D[i, :] = np.sqrt(np.sum((X - row)**2, 1))

%timeit pairwise_numpy_fast(X, D)

10 loops, best of 3: 46.9 ms per loop

That’s a much more satisfying result. Goes to show how much one can gain without even going too crazy with tricks.

Then there is my current approach, using Fortran for the tight loops. Notice how much the Fortran source resembles the Python ones. This is not your daddy’s Fortran, it’s Fortran 90.

subroutine pairwise(X, M, N, D)
  double precision, intent (IN) :: X(M, N)
  integer, intent(IN) :: M, N
  double precision, intent (OUT) :: D(M, M)
  integer :: i, j

  forall (i = 1:M, j = 1:M)
     D(i, j) = sqrt(sum((X(i, :) - X(j, :))**2))
  end forall
end subroutine pairwise

After compiling this file to a Python module with f2py (

f2py -c -m pairwise_fortran pairwise.f90

), we can check it with

from pairwise_fortran import pairwise as pairwise_fortran
%timeit pairwise_fortran(X, M, N)

100 loops, best of 3: 6.58 ms per loop

Now that’s more like it!

So, from 8.3 s down to 6.6 ms. That’s three orders of magnitude from the original naive Python code, but the “clever” numpy version is less than one order slower. My feeling is that the benefits gained from a language like Python more than make up for any speed deficits. Of course, there’s the distinct advantage of having a really quick way to incorporate the efficient compilers available for Fortran using a tool like f2py.