Applying Minimum Image Convention in Python

Question

I am computing pairwise Euclidean distances for 3-D vectors representing particle positions in a periodic system. The minimum image convention is applied for each periodic boundary such that a particle only considers the nearest image of another particle when computing the distance.

This code is part of a larger post-processing effort in Python. Here is a working example of what I am implementing:

import numpy as np
from scipy.spatial.distance import cdist

s = np.array([40, 30, 20])
half_s = 0.5*s

a = np.transpose(np.random.rand(1000,3) * s)
b = np.transpose(np.random.rand(1000,3) * s)

dists = np.empty((3, a.shape[1]*b.shape[1]))
for i in range(3):
    dists[i,:] = cdist(a[i,:].reshape(-1,1), b[i,:].reshape(-1,1), 'cityblock').ravel()
    dists[i,:] = np.where(dists[i,:] > half_s[i], dists[i,:] - s[i], dists[i,:])
dists = np.sqrt(np.einsum("ij,ij->j", dists, dists))

The domain size s and the 3xn particle position arrays a and b are obtained from existing data structures, but this example uses sizes I would typically expect. I should emphasize that a and b can have different lengths, but on average I expect the final dists array to represent around a million distances (+/- an order of magnitude).

The last 5 lines computing the distances will need to be run many thousands of times, so this is what I hope to optimize.

The difficulty arises from the need to apply the minimum image convention to each component independently. I haven't been able to find anything which can beat SciPy's cdist for computing the unsigned distance components, and NumPy's einsum function seems to be the most efficient way to reduce the distances for arrays of this size. At this point, the bottleneck in the code is in the NumPy where function. I also tried using NumPy's casting method by replacing the penultimate line with

dists[i,:] -= (dists[i,:] * s_r[i] + 0.5).astype(int) * s[i]

where s_r = 1/s, but this yielded the same runtime. This paper discusses various techniques to handle this operation in C/C++, but I'm not familiar enough with the underlying CPython implementation in NumPy/SciPy to determine what's best here. I'd love to parallelize this section, but I had little success with the multiprocessing module and cdist is incompatible with Numba. I would also entertain suggestions on writing a C extension, though I've never incorporated one in Python before.

Answer 1

I think you could move the where function out of the loop:

dists = np.empty((3, a.shape[1] * b.shape[1]))

for i in range(3):
    dists[i, :] = cdist(a[i, :].reshape(-1, 1),
                        b[i, :].reshape(-1, 1), 'cityblock').reshape(-1)

dists = np.where(dists > half_s[..., None], dists - s[..., None], dists)
...

Answer 2

def distance(coords1, coords2, lattice_matrix, pairwise=False):
"""
Calculates the distance between two arrays of points in 3D space under
periodic boundary conditions.

Parameters:
coords1 (array-like): the coordinates of the first point
coords2 (array-like): the coordinates of the second point
lattice_matrix (array-like): the 3x3 lattice matrix of the unit cell
pairwise (bool): if True, returns matrix of pairwise distances between all 
  points in coords1 to all points in coords2 (slower)

Returns:
float: the distance between the two points under periodic boundary
  conditions
"""
if pairwise:
    diff = coords2[:, np.newaxis] - coords1[np.newaxis, :]
else:
    diff = coords2 - coords1

diff = np.dot(diff, np.linalg.inv(lattice_matrix))
diff -= np.round(diff)  # minimum image convention
diff = np.dot(diff, lattice_matrix)

return np.linalg.norm(diff, axis=-1)