Compute square distances from numpy array

Question

I'm having a brain lapse, but I just can't get this to work. I have an array of distances:

import numpy as np
zvals = np.linspace(-5,5,10)
d = np.array([(0,0,z) for z in zvals])

I want to compute the square distance of the points in the array. The non-numpy way to make this work is:

d2 = np.array([np.dot(d[i,:],d[i,:]) for i in range(d.shape[0])])

However, I know that there must be some way to do this with just a single call to dot, right? That being said, neither

d2 = np.dot(d,d.T)

or

d2 = np.dot(d.T,d)

give what I want. I'm being stupid, I realize, but please enlighten me here. Thanks!

Answer 1

Edit: As of NumPy 1.9, it appears inner1d may be faster. (Thanks to Nuno Aniceto for pointing this out):

In [9]:  %timeit -n 1000000 inner1d(d,d)
1000000 loops, best of 3: 1.39 µs per loop

In [14]: %timeit -n 1000000 einsum('ij,ij -> i', d, d)
1000000 loops, best of 3: 1.8 µs per loop

PS. Always test benchmarks for yourself on inputs similar to your intended use case. Results may vary for a variety of reasons, such as size of input, hardware, OS, Python version, NumPy version, compiler, and libraries (e.g. ATLAS, MKL, BLAS).

---

If you have NumPy version 1.6 or better, you could use np.einsum:

In [40]: %timeit np.einsum('ij,ij -> i', d, d)
1000000 loops, best of 3: 1.79 us per loop

---

In [46]: from numpy.core.umath_tests import inner1d

In [48]: %timeit inner1d(d, d)
100000 loops, best of 3: 1.97 us per loop

---

In [44]: %timeit np.sum(d*d, axis=1)
100000 loops, best of 3: 5.39 us per loop

---

In [41]: %timeit np.diag(np.dot(d,d.T)) 
100000 loops, best of 3: 7.2 us per loop

---

In [42]: %timeit np.array([np.dot(d[i,:],d[i,:]) for i in range(d.shape[0])])
10000 loops, best of 3: 26.1 us per loop

Answer 2

Dot product functions are very fast, and for really simple stuff may even beat np.einsum (which is a terrific function you should definitely learn to use). Numpy has a hidden little gem, inner1d, which does the dot product, with broadcasting, over the last dimensions of its arguments. You can use it as follows:

from numpy.core.umath_tests import inner1d
inner1d(a, a)