All-pairs Euclidean distance

Given two arrays of $d$-dimensional vectors, $\mathbf{A}\in\mathbb{R}^{M\times d}$ and $\mathbf{B}\in\mathbb{R}^{N\times d}$, we are interested in efficient vectorized code for obtaining a matrix $\mathbf{D}\in\mathbb{R}^{M\times N}$ such that

where $(\mathbf{D})_{ij}$ is the $ij$th entry of $\mathbf{D}$, $\mathbf{a}_i$ is the $i$th row of $\mathbf{A}$ and $\mathbf{b}_j$ is the $j$th row of $\mathbf{B}$. In other words, we want to compute the Euclidean distance between all vectors in $\mathbf{A}$ and all vectors in $\mathbf{B}$. The following numpy code does exactly this:

def all_pairs_euclid_naive(A, B):
	#
	D = numpy.zeros((A.shape[0], B.shape[0]), dtype=numpy.float32)
	for i in range(0, D.shape[0]):
		for j in range(0, D.shape[1]):
			D[i, j] = numpy.linalg.norm(A[i, :] - B[j, :])
	#
	return D

Unfortunately, this code is really inefficient. To rectify the issue, we need to write a vectorized version in which we avoid the explicit usage of loops. To arrive at a solution, we first expand the formula for the Euclidean distance:

This leads us to the following equation for $\mathbf{D}$:

where $\mathbf{S}_A\in\mathbb{R}^{M\times N}$ is such that $(\mathbf{S}_A)_{ij}=\mathbf{a}_i^T\mathbf{a}_i$ and $\mathbf{S}_B\in\mathbb{R}^{M\times N}$ is such that $(\mathbf{S}_B)_{ij}=\mathbf{b}_j^T\mathbf{b}_j$. The square root is taken elementwise.

Each of the three above terms can be obtain with vectorized code.

The first term is obtained for each $i$ by squaring the entries of $\mathbf{A}$, summing along the second dimension after that and repeating the obtained result $N$ times to obtain an $M\times N$ matrix. The second term can be computed with the standard matrix-matrix multiplication routine. The third term is obtained in a simmilar manner to the first term.

Without further ado, here is the numpy code:

def all_pairs_euclid_numpy(A, B):
	#
	sqrA = numpy.broadcast_to(numpy.sum(numpy.power(A, 2), 1).reshape(A.shape[0], 1), (A.shape[0], B.shape[0]))
	sqrB = numpy.broadcast_to(numpy.sum(numpy.power(B, 2), 1).reshape(B.shape[0], 1), (B.shape[0], A.shape[0])).transpose()
	#
	return numpy.sqrt(
		sqrA - 2*numpy.matmul(A, B.transpose()) + sqrB
	)

And the following implementation uses Pytorch:

def all_pairs_euclid_torch(A, B):
	#
	sqrA = torch.sum(torch.pow(A, 2), 1, keepdim=True).expand(A.shape[0], B.shape[0])
	sqrB = torch.sum(torch.pow(B, 2), 1, keepdim=True).expand(B.shape[0], A.shape[0]).t()
	#
	return torch.sqrt(
		sqrA - 2*torch.mm(A, B.t()) + sqrB
	)

Note that the above code can be executed on the GPU as well.

By setting $M=2048$, $N=2048$ and $d=2048$, we obtain the following timings for the modern 40-thread Intel CPU and the Nvidia 1080Ti GPU:

You can get the time measurements for your hardware by running the following script.