Поиск индекса стоимости в массивах 3D

Question

import numpy as np
# The 3D arrays have the axis: Z, X, Y
arr_keys = np.random.rand(20, 5, 5)
arr_vals = np.random.rand(20, 5, 5)
arr_idx = np.random.rand(5, 5)

Для каждой ячейки сетки в arr_idx я хочу посмотреть Z-позицию ближайшего к ней значения в arr_keys (но с тем же положением X, Y) и вернуть значение в соответствующей позиции в arr_vals массив. Есть ли способ сделать это без использования вложенных циклов?

Итак, если значение в X = 0, Y = 0 для arr_idx равно 0,5, я хочу найти ближайший к нему номер в X = 0, Y = 0, Z колеблется от 0 до 10 в arr_keys, а затем я хочу использовать позицию Z этого числа (давайте назовем его Z_prime), чтобы найти значение в arr_vals (Z_prime, X = 0, Y = 0)

Answer 1

Это тип проблемы, для которой было создано np.take_along_axis:

# shape (20, 5, 5)
diff = np.abs(arr_idx - arr_keys)

# argmin(..., keepdims=True) doesn't exist yet - this emulates it
# shape (1, 5, 5)
inds = np.expand_dims(np.argmin(diff, axis=0), axis=0)

# shape (1, 5, 5)
res = np.take_along_axis(arr_vals, inds, axis=0)

# shape (5, 5)
res = res.squeeze(axis=0)

Answer 2

Я думаю, что ответ @ xnx довольно хороший. У меня больше, но я все равно выложу;).

Также обратите внимание: NumPy создан для эффективной обработки больших многомерных массивов путем векторизации операций. Поэтому я бы рекомендовал избегать циклов for как можно больше. Какую бы задачу вы ни искали, есть (обычно) способ сделать это, избегая петель.

arr_keys = np.split(arr_keys, 20)
arr_keys = np.stack(arr_keys, axis=-1)[0]
arr_vals = np.split(arr_vals, 20)
arr_vals = np.stack(arr_vals, axis=-1)[0]

arr_idx = np.expand_dims(arr_idx, axis=-1)

difference = np.abs(arr_keys - arr_idx)
minimum = np.argmin(difference, axis=-1)

result = np.take_along_axis(arr_vals, np.expand_dims(minimum, axis=-1), axis=-1)
result = np.squeeze(result, axis=-1)

Answer 3

Немного многословно, чем уже опубликованное решение, но легче для понимания.

import numpy as np
# The 3D arrays have the axis: Z, X, Y
arr_keys = np.random.rand(20, 5, 5)
arr_vals = np.random.rand(20, 5, 5)
arr_idx = np.random.rand(5, 5)

arr_idx = arr_idx[np.newaxis, :, :]
dist = np.abs(arr_idx - arr_keys)
dist_ind = np.argmin(dist, axis=0)
x = np.arange(0, 5, 1)
y = np.arange(0, 5, 1)
xx, yy = np.meshgrid(x, y)

res = arr_vals[dist_ind, yy, xx]

Answer 4

Я думаю, что это может сработать: сверните оси в правильной ориентации, найдите индекс значения (абсолютного) минимума для каждого из значений 5x5 X, Y и возьмите соответствующие Z-значения из arr_vals:

idx = np.argmin(np.abs(np.rollaxis(arr_keys,0,3) - arr_idx[:,:,None]), axis=2)
i,j = np.ogrid[:5,:5]
arr_vals[idx[i,j],i,j]

Чтобы проверить это, попробуйте случай (3,2,2):

In [15]: arr_keys
Out[15]: 
array([[[ 0.19681533,  0.26897784],
        [ 0.60469711,  0.09273087]],

       [[ 0.04961604,  0.3460404 ],
        [ 0.88406912,  0.41284309]],

       [[ 0.46298201,  0.33809574],
        [ 0.99604152,  0.4836324 ]]])

In [16]: arr_vals
Out[16]: 
array([[[ 0.88865681,  0.88287688],
        [ 0.3128103 ,  0.24188022]],

       [[ 0.23947227,  0.57913325],
        [ 0.85768064,  0.91701097]],

       [[ 0.78105669,  0.84144339],
        [ 0.81071981,  0.69217687]]])

In [17]: arr_idx
Out[17]: 
array([[[ 0.31352609],
        [ 0.75462329]],

       [[ 0.44445286],
        [ 0.97086161]]])

дает:

array([[ 0.88865681,  0.57913325],
       [ 0.3128103 ,  0.69217687]])