pairwise distance calculation for n-dimensional coordinates

pdist_pbc.py

import numpy as np
from scipy.spatial.distance import pdist, squareform
import matplotlib.pyplot as plt


def pdist_pbc(positions, box):
    """
    Get the pair-wise distances of particles in a priodic boundary box

    Args:
        positiosn (:obj:`numpy.ndarray`): coordinate of particles, shape (N, dim)
        box (:obj:`float`): the length of the priodic bondary.
            The box should be cubic

    Return:
        the pairwise distance, shape ( (N * N - N) / 2, ), use
            :obj:`scipy.spatial.distance.squareform` to recover the matrix form.
    """
    n, dim = positions.shape
    result = np.zeros(int((n * n - n) / 2), dtype=np.float64)
    for d in range(dim):
        dist_1d = pdist(positions[:, d][:, np.newaxis])
        dist_1d[dist_1d > box / 2] -= box
        result += np.power(dist_1d, 2)
    return np.sqrt(result)


if __name__ == "__main__":
    dists = pdist_pbc(np.random.random((1000, 3)), box=1)
    plt.hist(dists, bins=200, density=True)
    plt.plot([np.sqrt(3)/2, np.sqrt(3)/2], [0, 5], color='k', ls='--')
    plt.ylim(0, 4)
    plt.show()