Optimizing Python distance calculation while accounting for periodic boundary conditions

Question

I have written a Python script to calculate the distance between two points in 3D space while accounting for periodic boundary conditions. The problem is that I need to do t

Answer 1

You should write your distance() function in a way that you can vectorise the loop over the 5711 points. The following implementation accepts an array of points as either the x0 or x1 parameter:

def distance(x0, x1, dimensions):
    delta = numpy.abs(x0 - x1)
    delta = numpy.where(delta > 0.5 * dimensions, delta - dimensions, delta)
    return numpy.sqrt((delta ** 2).sum(axis=-1))

Example:

>>> dimensions = numpy.array([3.0, 4.0, 5.0])
>>> points = numpy.array([[2.7, 1.5, 4.3], [1.2, 0.3, 4.2]])
>>> distance(points, [1.5, 2.0, 2.5], dimensions)
array([ 2.22036033,  2.42280829])

The result is the array of distances between the points passed as second parameter to distance() and each point in points.