find pairwise row differences of numpy array (with perdioic boundaries)

Question

I have a (N,3)-dimensional numpy array holding the 3-dim. positions of N particles in its rows. I want to obtain the pairwise distances of these particles using periodic boundaries (the latter means that if particles have a x-, y-, or z-distance d that is larger than a value L/2, then we take the distance abs(d-L) rather than d itself).

My C-like prototype looks like this:

def get_distances(positions, L):
    """Compute distances from positions using periodic boundaries"""
    
    rs = []
    L_half = 0.5*L
    
    for pos_i in positions:
        for pos_j in positions:
            rx = np.abs(pos_i[0] - pos_j[0])
            rx = rx if rx < L_half else np.abs(rx-L)

            ry = np.abs(pos_i[1] - pos_j[1])
            ry = ry if ry < L_half else np.abs(ry-L)

            rz = np.abs(pos_i[2] - pos_j[2])
            rz = rz if rz < L_half else np.abs(rz-L)            
        
            rs += [np.sqrt(rx*rx + ry*ry + rz*rz)]
    
    return rs

How can I make this more efficient and Pythonic?

Answer 1

Are you sure your provided pseudocode is correct? In particular, the line rs += [np.sqrt(rx*rx + ry*ry + rz*rz)] is in the outer loop, so it would only use the rx, ry, rz computed in the last iteration of the inner loop.

These are some numpy tricks that you might find useful for your problem:

import numpy as np

N, L = 10, 1
arr = np.random.rand(N, 3)

# You can use broadcasting to do an operation between
# each pair of elements within an array:
r = np.abs(arr[:, None, :] - arr[None, :, :])
# `r.shape == (N, N, 3)`; the absolute difference
# between each pair of shape (3,) subvectors in `arr`.

# if-else logic is solved using `np.where`:
s = np.where(r < L/2, r, np.abs(r - L))

# np.sqrt(x**2 + y**2 + ...) is done with `np.linalg.norm`:
out = np.linalg.norm(s, axis=2)

Which leaves you with a shape (N, N) array!