Эффективная реализация Python сравнений массивов numpy

фон

у меня есть два массива numpy, которые я хотел бы использовать для выполнения некоторых операций сравнения наиболее эффективным/быстрым способом. Оба содержат только неподписанные ints.

pairs - это n x 2 x 3 массив, который содержит длинный список парных 3D-координат (для некоторой номенклатуры pairs массив содержит набор пар...) - т. е.

# full pairs array
In [145]: pairs
Out[145]:
    array([[[1, 2, 4],
        [3, 4, 4]],
        .....
       [[1, 2, 5],
        [5, 6, 5]]])

# each entry contains a pair of 3D coordinates
In [149]: pairs[0]
Out[149]:
array([[1, 2, 4],
       [3, 4, 4]])

positions это n x 3 массив, который содержит набор 3D координаты

In [162]: positions
Out[162]:
array([[ 1,  2,  4],
       [ 3,  4,  5],
       [ 5,  6,  3],
       [ 3,  5,  6],
       [ 6,  7,  5],
       [12,  2,  5]])

цель Я хочу создать массив, который является подмножеством pairs массив, но содержит только записи, где не более одной из пар находится в массиве позиций-т. е. не должно быть пар, где обе пары лежат в массиве позиций. Для некоторой информации о домене каждая пара будет иметь по крайней мере одну из парных позиций в списке позиций.

подходы пытались до сих пор Мой первоначальный наивный подход к петля над каждой парой в pairs массив и вычесть каждую из двух парных позиций из positions вектор, определяющий, если в обоих случаях мы нашли совпадение, указанное наличием 0 в обоих векторах, которые исходят из операций вычитания:

 if (~(positions-pair[0]).any(axis=1)).any() and 
    (~(positions-pair[1]).any(axis=1)).any():
    # both members of the pair were in the positions array -
    # these weren't the droids we were looking for
    pass
 else:
    # append this set of pairs to a new matrix

это отлично работает и используеткакой-то векторизация, но, вероятно, есть лучший способ сделать это?

для некоторых других чувствительных к производительности частей этой программы у меня есть переписанные вещи с Cython, которые привели к массовому ускорению, хотя в этом случае (по крайней мере, на основе наивной вложенной реализации для цикла) это было немного медленнее, чем подход, описанный выше.

если у людей есть предложения, я рад профилировать и отчитываться (у меня есть вся инфраструктура профилирования).

7
cython numpy performance python

2 ответов

подход #1

как уже упоминалось в вопросе, оба массива содержат только unsigned ints, который может быть использован для слияния XYZв эквивалентную версию линейных индексов, которая была бы уникальной для каждого уникального XYZ триплет. Реализация будет выглядеть примерно так -

maxlen = np.max(pairs,axis=(0,1))
dims = np.append(maxlen[::-1][:-1].cumprod()[::-1],1)

pairs1D = np.dot(pairs.reshape(-1,3),dims)
positions1D = np.dot(positions,dims)
mask_idx = ~(np.in1d(pairs1D,positions1D).reshape(-1,2).all(1))
out = pairs[mask_idx]

поскольку вы имеете дело с 3D-координатами, вы также можете использоватьcdist для проверки идентичны XYZ тройни между входной массив. Далее перечислены две реализации с учетом этой идеи.

подход #2

from scipy.spatial.distance import cdist

p0 = cdist(pairs[:,0,:],positions)
p1 = cdist(pairs[:,1,:],positions)
out = pairs[((p0==0) | (p1==0)).sum(1)!=2]

подход #3

mask_idx = ~((cdist(pairs.reshape(-1,3),positions)==0).any(1).reshape(-1,2).all(1))
out = pairs[mask_idx]

во время выполнения тестов -

In [80]: n = 5000
    ...: pairs = np.random.randint(0,100,(n,2,3))
    ...: positions= np.random.randint(0,100,(n,3))
    ...: 

In [81]: def cdist_split(pairs,positions):
    ...:    p0 = cdist(pairs[:,0,:],positions)
    ...:    p1 = cdist(pairs[:,1,:],positions)
    ...:    return pairs[((p0==0) | (p1==0)).sum(1)!=2]
    ...: 
    ...: def cdist_merged(pairs,positions):
    ...:    mask_idx = ~((cdist(pairs.reshape(-1,3),positions)==0).any(1).reshape(-1,2).all(1))
    ...:    return pairs[mask_idx]
    ...: 
    ...: def XYZ_merged(pairs,positions):
    ...:    maxlen = np.max(pairs,axis=(0,1))
    ...:    dims = np.append(maxlen[::-1][:-1].cumprod()[::-1],1)
    ...:    pairs1D = np.dot(pairs.reshape(-1,3),dims)
    ...:    positions1D = np.dot(positions,dims)
    ...:    mask_idx1 = ~(np.in1d(pairs1D,positions1D).reshape(-1,2).all(1))
    ...:    return pairs[mask_idx1]
    ...: 

In [82]: %timeit cdist_split(pairs,positions)
1 loops, best of 3: 662 ms per loop

In [83]: %timeit cdist_merged(pairs,positions)
1 loops, best of 3: 615 ms per loop

In [84]: %timeit XYZ_merged(pairs,positions)
100 loops, best of 3: 4.02 ms per loop

проверить результаты -

In [85]: np.allclose(cdist_split(pairs,positions),cdist_merged(pairs,positions))
Out[85]: True

In [86]: np.allclose(cdist_split(pairs,positions),XYZ_merged(pairs,positions))
Out[86]: True
6

уточняя мой комментарий:

расширения pairs чтобы быть более интересным. Не стесняйтесь тестировать с большим, более реалистичным массивом:

In [260]: pairs = np.array([[[1,2,4],[3,4,4]],[[1,2,5],[5,6,5]],[[3,4,5],[3,5,6]],[[6,7,5],[1,2,3]]])

In [261]: positions = np.array([[ 1,  2,  4],
       [ 3,  4,  5],
       [ 5,  6,  3],
       [ 3,  5,  6],
       [ 6,  7,  5],
       [12,  2,  5]])

разверните оба массива в широковещательные формы:

In [262]: I = pairs[None,...]==positions[:,None,None,:]

In [263]: I.shape
Out[263]: (6, 4, 2, 3)

большой логический массив, показывающий соответствие элементов по всем измерениям. Fell свободно заменять другие сравнения (difference ==0, np.isclose для поплавков и т. д.). 

In [264]: J = I.all(axis=-1).any(axis=0).sum(axis=-1)

In [265]: J
Out[265]: array([1, 0, 2, 1])

консолидировать результаты по различным измерениям. Матч все числа по координатам, совпадают с любыми позициями, подсчет совпадений по парам.

In [266]: pairs[J==1,...]
Out[266]: 
array([[[1, 2, 4],
        [3, 4, 4]],

       [[6, 7, 5],
        [1, 2, 3]]])

J==1 представляют элементы, где совпадает только одно значение пары. (см. примечание в конце)

сочетание any, and и sum работа это для теста, но может потребоваться регулировка с большим тестом(ы). Но идея вообще применима.

для размера массивов, которые https://stackoverflow.com/a/31901675/901925 тесты, мое решение довольно медленно. Особенно это делает == тест, который приводит к I С (5000, 5000, 2, 3).

сжатие последнего измерения очень помогает

dims = np.array([10000,100,1])  # simpler version of dims from XYZmerged
pairs1D = np.dot(pairs.reshape(-1,3),dims)
positions1D = np.dot(positions,dims)
I1d = pairs1D[:,None]==positions1D[None,:]
J1d = I1d.any(axis=-1).reshape(pairs.shape[:2]).sum(axis=-1)

я изменил J1d expression to match mine-подсчитать количество совпадений на пару.

на in1d1 это Divakar использует еще быстрее:

mask = np.in1d(pairs1D, positions1D).reshape(-1,2)
Jmask = mask.sum(axis=-1)

я только что понял, что ОП просит at most one of the pairs is in the positions array. Где, как я проверяю для exactly one match per pair. поэтому все мои тесты должны быть изменены на pairs[J<2,...].

(в моей конкретной случайной выборке для n=5000 это оказывается всем. Нет pairs где оба находятся в positions. 54 из 5000 у J==1, остальные 0-Нет).

3