[Python] Fastest way to find the nearest pairs between two numpy arrays without duplicates

Given two large numpy arrays A and B with different number of rows (len(B) > len(A)) but same number of columns (A.shape[1] = B.shape[1] = 3). I want to know the fastest way to get a subset C from B that has the minimum total distance (sum of all pair-wise distances) to A without duplicates (each pair must be both unique). This means C should have the same shape as A.

Below is my code, but there are two main issues:

1. I cannot tell if this gives the minimum total distance
2. In reality I have a much more expensive distance-calculating function rather than np.linalg.norm (needs to take care of periodic boundary conditions). I think this is definitely not the fastest way to go since the code below calls the distance-calculating function one pair per time. There is a significant overhead when I call the more expensive distance-calculating function and it will run forever. Any suggestions?

import numpy as np
from operator import itemgetter
import random
import time

A = 100.*np.random.rand(1000, 3)
B = A.copy()
for (i,j), _ in np.ndenumerate(B):
B[i,j] += np.random.rand()
B = np.vstack([B, 100.*np.random.rand(500, 3)])

def calc_dist(x, y):
return np.linalg.norm(x - y)

t0 = time.time()
taken = []
for rowi in A:
res = min(((k, calc_dist(rowi, rowj)) for k, rowj in enumerate(B)
if k not in taken), key=itemgetter(1))
taken.append(res[0])

C = B[taken]

print(A.shape, B.shape, C.shape)
>>> (1000, 3) (1500, 3) (1000, 3)

print(time.time() - t0)
>>> 12.406389951705933


Edit: for those who are interested in the expensive distance-calculating function, it uses the ase package (can be installed by pip install ase)

from ase.geometry import find_mic
def calc_mic_dist(x, y):
return find_mic(np.array([x]) - np.array([y]),
cell=np.array([[50., 0.0, 0.0],
[25., 45., 0.0],
[0.0, 0.0, 100.]]))[1][0]

Continue reading...