Самый быстрый способ найти ближайшие пары между двумя массивами numpy без дубликатов
Даны два больших массива numpy A и B с разным количеством строк (len(B) > len(A)), но одинаковым количеством столбцов (A.shape[1] = B.shape[1] = 3). Я хочу знать самый быстрый способ получить подмножество C из B, которое имеет минимальное общее расстояние (сумма всех попарных расстояний) до A без дубликатов (каждая пара должна быть уникальной). Это означает, что C должен иметь ту же форму, что и A.
Ниже мой код, но есть две основные проблемы:
- Я не могу сказать, дает ли это минимальное общее расстояние
- На самом деле у меня гораздо более дорогая функция вычисления расстояния, чем
np.linalg.norm(необходимо учитывать периодические граничные условия). Я думаю, что это определенно не самый быстрый способ, поскольку приведенный ниже код вызывает функцию расчета расстояния по одной паре за раз. Когда я вызываю более затратную функцию расчета расстояния, возникают значительные накладные расходы, и она будет работать вечно. Какие-либо предложения?
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
Редактировать: для тех, кто интересуется дорогой функцией расчета расстояния, он использует пакет ase (может быть установлен 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]
Я добавляю дорогостоящую функцию расчета расстояния.
Просто мысль, как насчет того, чтобы отсортировать массивы A и B и начать рассматривать элементы в B для подмножества C с n-го элемента в B, который ближе всего к первому элементу в A? (дополнительная деталь: если значение n больше len(B) - len(A), то n можно сдвинуть назад в качестве компромисса) Я думаю, что таким образом будут генерироваться лучшие пары и обеспечиваться наименьшее расстояние. Имеет ли это смысл?
Да, я согласен с тем, что A и B должны быть сначала отсортированы каким-то образом, чтобы каждая пара всегда была «идеальной» парой, но я не совсем понимаю часть сдвига назад. Было бы более понятно, если бы вы могли написать какой-нибудь псевдокод.
Ответы 2
Использование мощности вещания и векторизации numpy
find_mic метод в ase.geometry может обрабатывать массивы 2d np.
from ase.geometry import find_mic
def calc_mic_dist(x, y):
return find_mic(x - y,
cell=np.array([[50., 0.0, 0.0],
[25., 45., 0.0],
[0.0, 0.0, 100.]]))[1]
Тест:
x = np.random.randn(1,3)
y = np.random.randn(5,3)
print (calc_mic_dist(x,y).shape)
# It is a distance metrics so:
assert np.allclose(calc_mic_dist(x,y), calc_mic_dist(y,x))
Выход:
(5,)
Как видите, метрики рассчитываются для каждого значения x с каждым значением y, потому что x-y в numpy выполняет магию трансляции.
Решение:
def calc_mic_dist(x, y):
return find_mic(x - y,
cell=np.array([[50., 0.0, 0.0],
[25., 45., 0.0],
[0.0, 0.0, 100.]]))[1]
t0 = time.time()
A = 100.*np.random.rand(1000, 3)
B = 100.*np.random.rand(5000, 3)
selected = [np.argmin(calc_mic_dist(a, B)) for a in A]
C = B[selected]
print (A.shape, B.shape, C.shape)
print (f"Time: {time.time()-t0}")
Выход:
(1000, 3) (5000, 3) (1000, 3)
Time: 9.817562341690063
Занимает около 10 секунд на коллабе Google
Тестирование:
Мы знаем, что calc_mic_dist(x,x) == 0, поэтому если A является подмножеством B, то C должно быть точно A
A = 100.*np.random.rand(1000, 3)
B = np.vstack([100.*np.random.rand(500, 3), A, 100.*np.random.rand(500, 3)])
selected = [np.argmin(calc_mic_dist(a, B)) for a in A]
C = B[selected]
print (A.shape, B.shape, C.shape)
print (np.allclose(A,C))
Выход:
(1000, 3) (2000, 3) (1000, 3)
True
Редактировать 1: Избегайте дубликатов
После того, как вектор в
Bвыбран, его нельзя снова выбрать для других значенияA
Этого можно добиться, удалив выбранный вектор из B после его выбора, чтобы он больше не появлялся в следующих строках A в качестве возможного кандидата.
A = 100.*np.random.rand(1000, 3)
B = np.vstack([100.*np.random.rand(500, 3), A, 100.*np.random.rand(500, 3)])
B_ = B.copy()
C = np.zeros_like(A)
for i, a in enumerate(A):
s = np.argmin(calc_mic_dist(a, B_))
C[i] = B_[s]
# Remove the paried
B_ = np.delete(B_, (s), axis=0)
print (A.shape, B.shape, C.shape)
print (np.allclose(A,C))
Выход:
(1000, 3) (2000, 3) (1000, 3)
True
Вещание numpy — это хорошо, но ваш код будет генерировать дубликаты. Например, одна точка в A может быть ближайшей точкой к нескольким точкам в B, но я хочу добиться того, чтобы после того, как точка в A образовала пару, ее больше нельзя было соединить снова.
@ShaunHan Надеюсь, редактирование удовлетворит ваши требования.
Спасибо. Я думаю, что это близко, но все же не совсем то, что я хочу, потому что это не может гарантировать минимальное общее расстояние. Я думаю, что A и B должны быть сначала отсортированы каким-то образом, чтобы каждая пара была «идеальной» парой.
Если вы согласны с расчетом всех расстояний N², что не так дорого для указанных вами размеров, у scipy.optimize есть функция, которая решит эту проблему напрямую.
import scipy.optimize
cost = np.linalg.norm(A[:, np.newaxis, :] - B, axis=2)
_, indexes = scipy.optimize.linear_sum_assignment(cost)
C = B[indexes]
Возможно, вы могли бы векторизовать вычисление функции, вы можете описать это?