Different ways to get pairwise distance in cubic box with PBC in Python

pdist_pbc.py

"""
Different functions to get pairwise distance in cubic box with PBC

Args:
    positions: the positions of particles, shape (N, dim)
    box (int or float): the size of the box

Return:
    np.ndarray: the pair-wise distances, shape (N, N)
"""
import numpy as np
from scipy.spatial.distance import pdist, squareform

def get_pdist_pbc_v1(positions, box):
    pos_in_pbc = positions % box
    pair_distances = np.empty((N, N))
    for i, p1 in enumerate(pos_in_pbc):
        for j, p2 in enumerate(pos_in_pbc):
            dist_nd_sq = 0
            for d in range(dim):
                dist_1d = abs(p2[d] - p1[d])
                if dist_1d > (box / 2):  # check if d_ij > box_size / 2
                    dist_1d = box - dist_1d
                dist_nd_sq += dist_1d ** 2
            pair_distances[i, j] = dist_nd_sq
    return np.sqrt(pair_distances)
    

def get_pdist_pbc_v2(positions, box):
    pos_in_pbc = positions % box
    dist_nd_sq = np.zeros(N * (N - 1) // 2)
    for d in range(dim):
        pos_1d = pos_in_pbc[:, d][:, np.newaxis]
        dist_1d = pdist(pos_1d)
        dist_1d[dist_1d > box * 0.5] -= box
        dist_nd_sq += dist_1d ** 2
    dist_nd = squareform(np.sqrt(dist_nd_sq))
    return dist_nd
    
    
def get_pdist_pbc_v3(positions, box):
    pos_in_box = positions.T / box
    pair_shift = pos_in_box[:, :, np.newaxis] - pos_in_box[:, np.newaxis, :]
    pair_shift = pair_shift - np.rint(pair_shift)
    return np.linalg.norm(pair_shift, axis=0) * box
    

if __name__ == "__main__":
    from time import time
    
    N, dim, box = 1000, 3, 2
    positions = np.random.uniform(0, 1, (N, dim))
    t0 = time()
    pd1 = get_pdist_pbc_v1(positions, box)
    print(f"{time() - t0:.4f} s spent by method 1")
    t0 = time()
    pd2 = get_pdist_pbc_v2(positions, box)
    print(f"{time() - t0:.4f} s spent by method 2")
    t0 = time()
    pd3 = get_pdist_pbc_v3(positions, box)
    print(f"{time() - t0:.4f} s spent by method 3")
    
    eq12 = np.allclose(pd1, pd2)
    eq13 = np.allclose(pd1, pd3)
    print("Getting identical results?", eq12 and eq13)