Question

The numpy version makes use of advanced broadcasting. To follow the code below, we will have to understand numpy broadcasting rules a little better. Here is the gist from:

From http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#numpy.newaxis

When operating on two arrays, NumPy compares their shapes element-wise. It starts with the trailing dimensions, and works its way forward. Two dimensions are compatible when

• they are equal, or
• one of them is 1

Arrays do not need to have the same number of dimensions. When either of the dimensions compared is one, the larger of the two is used. In other words, the smaller of two axes is stretched or “copied” to match the other.

Distance between scalars

x = np.arange(10)
x

array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

# if we insert an extra dimension into x with np.newaxis
# we get a (10, 1) matrix
x[:, np.newaxis].shape

(10, 1)

Comparing shape

x[:, None] = 10 x 1
x          =     10

When we subtract the two arrays, broadcasting rules first match the the trailing axis to 10 (so x[:, None] is stretched to be (10,10)), and then matching the next axis, x is stretechd to also be (10,10).

# This is the pairwise distance matrix!
x[:, None] - x

array([[ 0, -1, -2, -3, -4, -5, -6, -7, -8, -9],
       [ 1,  0, -1, -2, -3, -4, -5, -6, -7, -8],
       [ 2,  1,  0, -1, -2, -3, -4, -5, -6, -7],
       [ 3,  2,  1,  0, -1, -2, -3, -4, -5, -6],
       [ 4,  3,  2,  1,  0, -1, -2, -3, -4, -5],
       [ 5,  4,  3,  2,  1,  0, -1, -2, -3, -4],
       [ 6,  5,  4,  3,  2,  1,  0, -1, -2, -3],
       [ 7,  6,  5,  4,  3,  2,  1,  0, -1, -2],
       [ 8,  7,  6,  5,  4,  3,  2,  1,  0, -1],
       [ 9,  8,  7,  6,  5,  4,  3,  2,  1,  0]])

Distance between vectors

# Suppose we have a collection of vectors of dimeniosn 2
# In the example below, there are 5 such 2-vectors
# We want to calculate the Euclidean distance
# for all pair-wise comparisons in a 5 x 5 matrix

x = np.arange(10).reshape(5,2)
print x.shape
print x

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

x[:, None, :].shape

(5, 1, 2)

Comparing shape

x[:, None, :] = 5 x 1 x 2
x          =        5 x 2

From the rules of broadcasting, we expect the result of subtraction to be a 5 x 5 x 2 array. To calculate Euclidean distance, we need to find the square root of the sum of squares for the 5 x 5 collection of 2-vectors.

delta = x[:, None, :] - x
pdist = np.sqrt((delta**2).sum(-1))
pdist

array([[  0.  ,   2.83,   5.66,   8.49,  11.31],
       [  2.83,   0.  ,   2.83,   5.66,   8.49],
       [  5.66,   2.83,   0.  ,   2.83,   5.66],
       [  8.49,   5.66,   2.83,   0.  ,   2.83],
       [ 11.31,   8.49,   5.66,   2.83,   0.  ]])

Finally, we come to the anti-climax - a one-liner function!

def pdist_numpy(xs):
    return np.sqrt(((xs[:,None,:] - xs)**2).sum(-1))

print pdist_numpy(A)
%timeit pdist_numpy(xs)

[[ 0.  5.]
 [ 5.  0.]]
10 loops, best of 3: 94.2 ms per loop