Быстрая реализация для окрестностей в NumPy тензор

Question

У меня есть 3-тензор A в numpy и индекс (i,j,k).Мне нужно поместить в один (возможно, numpy) массив arr линейный индекс «окрестностей» (i,j,k) (т. Е. Элементов (i+i1, j+j1, k+k1), где i1, j1, k2 = -1, 0, 1).Мне нужно сделать это для каждого индекса (i,j,k) из A, поэтому мой вопрос, есть короткий и быстрый способ сделать это?

Пример игрушки:

n = 10
dim = (10, 10, 10)
i, j, k = (2, 3, 4)
res = []
pm = [-1, 0, 1]
for i1 in range(3):
    for j1 in range(3):
        for k1 in range(3):
            # ravel_multi_index transforms a multi-index (i,j,k) to the corresponding linear index of a tensor of shape "dim"
            res.append(np.ravel_multi_index((i + pm[i1], j + pm[j1], k + pm[k1]), dim))
#print...
[123, 124, 125, 133, 134, 135, 143, 144, 145, 223, 224, 225, 233, 234, 235, 243, 244, 245, 323, 324, 325, 333, 334, 335, 343, 344, 345]

Существует 27 окрестностей.

Индекс границы имеет меньше окрестностей: например, координата (0, 0, 0) имеет 8 окрестностей.

В качестве решения я думал о некоторых ""Соседство" соседств до одного исправлено, но это быстро становится громоздким.

Answer 1

Это должно сделать это.Вы просто генерируете массив vertices один раз и позволяете вещанию делать все остальное.vertices дает все 27 возможностей, необходимых для добавления в индекс и получения всех его соседей.

ijk = np.array([2, 3, 4]) # an example point 
dim = (10, 10, 10) # dimensions 

pts = np.stack(([-1,0,1],)*3,0)
vertices = (np.array(np.meshgrid(*pts)).T).reshape(3**3,3) # a 27x3 array
# the vertices are the points to add to ijk to get all neighbors 

n = (vertices + ijk).T # neighbors 
n = n.T[((n[0]>=0)&(n[0]<dim[0])
         &(n[1]>=0)&(n[1]<dim[1])
          &(n[2]>=0)&(n[2]<dim[2]))] # out of bounds elements filtered

Это возвращает

>>> n
array([[1, 2, 3],
       [1, 3, 3],
       [1, 4, 3],
       [2, 2, 3],
       [2, 3, 3],
       [2, 4, 3],
       [3, 2, 3],
       [3, 3, 3],
       [3, 4, 3],
       [1, 2, 4],
       [1, 3, 4],
       [1, 4, 4],
       [2, 2, 4],
       [2, 3, 4],
       [2, 4, 4],
       [3, 2, 4],
       [3, 3, 4],
       [3, 4, 4],
       [1, 2, 5],
       [1, 3, 5],
       [1, 4, 5],
       [2, 2, 5],
       [2, 3, 5],
       [2, 4, 5],
       [3, 2, 5],
       [3, 3, 5],
       [3, 4, 5]])

Вы можете принудительно ввести это в свой пример вывода с помощью

sorted([int(''.join(x)) for x in n.astype('str')])

, который возвращает

 [123,
 124,
 125,
 133,
 134,
 135,
 143,
 144,
 145,
 223,
 224,
 225,
 233,
 234,
 235,
 243,
 244,
 245,
 323,
 324,
 325,
 333,
 334,
 335,
 343,
 344,
 345]