Как эффективно вычислять парные количества в numpy?
Я хочу вычислить парные величины, например. расстояния между двумя точками.
Простым примером может быть
import numpy as np
N = 9
x = np.linspace(0,1,N)
y = np.abs(x - x[:,None]) # pairwise 1d eucdlidian distance
В результате получится массив (N,N), содержащий расстояния от каждого элемента в x до каждого другого элемента в x.
Однако нижний (или верхний) треугольник y можно отбросить, поскольку расстояние от x[0] до x[1] такое же, как расстояние от x[1] до x[0].
Может ли кто-нибудь придумать способ вообще никогда не вычислять нижний (или верхний) треугольник? В идеале это можно было бы обобщить на массивы Nd, например.
x = np.linspace(0,1,N)
y = x + (1j * x)[:,None] # complex plane
z = np.abs(y[:,:,None,None] - y[None,None,:,:]) # pairwise 2d euclidian distance
и любая попарная величина, например
y = np.random.randint(0,2,(N,N))
z = y[:,:,None,None] == y[None,None,:,:] # pairwise equality
Заранее спасибо.
PS: Для всех, кому интересно, это возникло, когда я попытался реализовать функцию энергии для 2D-моделирования отжига.
Ответы 1
Вручную
Вы можете сгенерировать пару индексов с помощью itertools.combinations и вручную вычислить расстояния:
from itertools import combinations
a, b = map(list, zip(*combinations(range(len(x)), 2)))
tri = abs(x[a]-x[b])
Выход:
array([0.125, 0.25 , 0.375, 0.5 , 0.625, 0.75 , 0.875, 1. , 0.125,
0.25 , 0.375, 0.5 , 0.625, 0.75 , 0.875, 0.125, 0.25 , 0.375,
0.5 , 0.625, 0.75 , 0.125, 0.25 , 0.375, 0.5 , 0.625, 0.125,
0.25 , 0.375, 0.5 , 0.125, 0.25 , 0.375, 0.125, 0.25 , 0.125])
Каковы расстояния между индексами:
[(0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7), (0, 8), (1, 2),
(1, 3), (1, 4), (1, 5), (1, 6), (1, 7), (1, 8), (2, 3), (2, 4), (2, 5),
(2, 6), (2, 7), (2, 8), (3, 4), (3, 5), (3, 6), (3, 7), (3, 8), (4, 5),
(4, 6), (4, 7), (4, 8), (5, 6), (5, 7), (5, 8), (6, 7), (6, 8), (7, 8)]
Использование scipy.spatial.distance.pdist
Именно это и делает pdist:
from scipy.spatial.distance import pdist, squareform
tri = pdist(x[:, None]) # one needs a 2D input
Выход:
array([0.125, 0.25 , 0.375, 0.5 , 0.625, 0.75 , 0.875, 1. , 0.125,
0.25 , 0.375, 0.5 , 0.625, 0.75 , 0.875, 0.125, 0.25 , 0.375,
0.5 , 0.625, 0.75 , 0.125, 0.25 , 0.375, 0.5 , 0.625, 0.125,
0.25 , 0.375, 0.5 , 0.125, 0.25 , 0.375, 0.125, 0.25 , 0.125])
Примечание. порядок пар такой же, как и в подходе itertools.combinations.
Если вы хотите получить квадратную форму:
squareform(tri)
array([[0. , 0.125, 0.25 , 0.375, 0.5 , 0.625, 0.75 , 0.875, 1. ],
[0.125, 0. , 0.125, 0.25 , 0.375, 0.5 , 0.625, 0.75 , 0.875],
[0.25 , 0.125, 0. , 0.125, 0.25 , 0.375, 0.5 , 0.625, 0.75 ],
[0.375, 0.25 , 0.125, 0. , 0.125, 0.25 , 0.375, 0.5 , 0.625],
[0.5 , 0.375, 0.25 , 0.125, 0. , 0.125, 0.25 , 0.375, 0.5 ],
[0.625, 0.5 , 0.375, 0.25 , 0.125, 0. , 0.125, 0.25 , 0.375],
[0.75 , 0.625, 0.5 , 0.375, 0.25 , 0.125, 0. , 0.125, 0.25 ],
[0.875, 0.75 , 0.625, 0.5 , 0.375, 0.25 , 0.125, 0. , 0.125],
[1. , 0.875, 0.75 , 0.625, 0.5 , 0.375, 0.25 , 0.125, 0. ]])
ND-массивы
Примечание. Это не распространяется непосредственно на массивы ND, но может временно преобразоваться в 2D в зависимости от конкретного варианта использования.
Во втором примере вы можете использовать:
# pairwise distances
tri = pdist(y.reshape(-1, 1))
# square form
out = squareform(tri).reshape(9, 9, 9, 9)
К сожалению, pdist, похоже, не обрабатывает сложные данные.
Ручная реализация с использованием itertools.combinations кажется универсальным, общим и простым в использовании решением, на которое я надеялся. Я проверю это и, если оно сработает, приму ваш ответ. Я бы предложил переместить его в начало вашего ответа, потому что другие решения более специфичны и требуют дополнительной установки.
Общий метод, предоставляемый numpy для этого, — ufunc.outer. Однако неясно, пропускает ли это обработку нижней треугольной части или даже ускоряет ее. Он также не поддерживает пользовательские вычисления, такие как евклидово расстояние. В качестве более общего метода вы можете использовать numba.