Как вырезать 3D-массив из 2D-массива без вложенных циклов?

Question

У меня есть четыре матрицы, частично созданные друг от друга:

• A - это трехмерная матрица, которая представляет собой стопку изображений в градациях серого и имеет форму (n, h, w)
• B - это трехмерная матрица, которая также представляет стопку изображений, где каждый срез индивидуально рассчитывается из соответствующего среза в A и также имеет форму (n, h, w)
• C - это двумерная матрица , содержащий индекс с максимальным значением B в направлении z и имеющий форму (h, w)
• D представляет собой 2D-матрицу, где значение из A копируется из определенного среза, который является обозначается значением в C в позиции (x, y).

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

import numpy as np

A = np.random.randint(0, 255, size=(3, 4, 5))
B = np.random.randint(0, 255, size=(3, 4, 5))
C = np.argmax(B, axis=0)

D = np.zeros(C.shape, dtype=int)
for y in range(C.shape[0]):
    for x in range(C.shape[1]):
        D[y, x] = A[C[y, x], y, x]


> A
array([[[ 24,  84, 155,   8, 147],
        [ 25,   4,  49, 195,  57],
        [ 93,  76, 233,  83, 177],
        [ 70, 211, 201, 132, 239]],

       [[177, 144, 247, 251, 207],
        [242, 148,  28,  40, 105],
        [181,  28, 132,  94, 196],
        [166, 121,  72,  14,  14]],

       [[ 55, 254, 140, 142,  14],
        [112,  28,  85, 112, 145],
        [ 16,  72,  16, 248, 179],
        [160, 235, 225,  14, 211]]])

> B
array([[[246,  14,  55, 163, 161],
        [  3, 152, 128, 104, 203],
        [ 43, 145,  59, 169, 242],
        [106, 169,  31, 222, 240]],

       [[ 41,  26, 239,  25,  65],
        [ 47, 252, 205, 210, 138],
        [194,  64, 135, 127, 101],
        [ 63, 208, 179, 137,  59]],

       [[112, 156, 183,  23, 253],
        [ 35,   6, 233,  42, 100],
        [ 66, 119, 102, 217,  64],
        [ 82,  67, 135,   6,   8]]])

> C
array([[0, 2, 1, 0, 2],
       [1, 1, 2, 1, 0],
       [1, 0, 1, 2, 0],
       [0, 1, 1, 0, 0]])

> D
array([[ 24, 254, 247,   8,  14],
       [242, 148,  85,  40,  57],
       [181,  76, 132, 248, 177],
       [ 70, 121,  72, 132, 239]])

Мой вопрос: как нарезать ломтиками A с C, эффективно устраняя вложенные циклы for? Моя первоначальная идея состояла в том, чтобы расширить C до трехмерной логической маски, где только позиции [c, y, x] установлены на True, а затем просто умножить ее поэлементно на A и взять сумму по оси z. Но я не могу представить себе реализацию Pythonesque без циклов (и мне, вероятно, больше не нужно было бы создавать булеву маску, если бы я знал, как это написать).

Ближайшая уже реализованная функция I найдено np.choose () , но для C требуется всего 32 элемента.

Answer 1

Стандартный подход здесь заключается в использовании np.take_along_axis() в сочетании с np.expand_dims() (основная идея представлена ​​также в документации np.argmax() ):

np.take_along_axis(A, np.expand_dims(C, axis=0), axis=0).squeeze()

Сравнивая предложенный подход с явными циклами и подходами на основе np.ogrid(), можно получить:

import numpy as np


def take_by_axis_loop_fix(arr, pos):
    result = np.zeros(pos.shape, dtype=int)
    for i in range(pos.shape[0]):
        for j in range(pos.shape[1]):
            result[i, j] = arr[pos[i, j], i, j]
    return result


def take_by_axis_ogrid_fix(arr, pos):
    i, j = np.ogrid[:pos.shape[0], :pos.shape[1]]
    return arr[pos[i, j], i, j]


def take_by_axis_np(arr, pos, axis=0):
    return np.take_along_axis(arr, np.expand_dims(pos, axis=axis), axis=axis).squeeze()


def take_by_axis_ogrid(arr, pos, axis=0):
    ij = tuple(np.ogrid[tuple(slice(None, d, None) for d in pos.shape)])
    ij = ij[:axis] + (pos[ij],) + ij[axis:]
    return arr[ij]

A_ = np.random.randint(0, 255, size=(300, 400, 500))
B_ = np.random.randint(0, 255, size=(300, 400, 500))
C_ = np.argmax(B_, axis=0)

funcs = take_by_axis_loop_fix, take_by_axis_ogrid_fix, take_by_axis_ogrid, take_by_axis_np
for func in funcs:
    print(func.__name__, np.all(func(A_, C_) == take_by_axis_loop_fix(A_, C_)))
    %timeit func(A_, C_)
    print()

# take_by_axis_loop_fix True
# 10 loops, best of 3: 114 ms per loop

# take_by_axis_ogrid_fix True
# 100 loops, best of 3: 5.94 ms per loop

# take_by_axis_ogrid True
# 100 loops, best of 3: 5.54 ms per loop

# take_by_axis_np True
# 100 loops, best of 3: 3.34 ms per loop

, указывающий, что это будет наиболее эффективным подход, предложенный до сих пор. Отметим также, что подходы на основе np.take_along_axis() и take_by_axis_ogrid() будут работать практически без изменений для входов с более высокой размерностью, в отличие от подходов _fix.

В частности, take_by_axis_ogrid() - это axis - agnosti c версия take_by_axis_ogrid_fix(), которая, по сути, nth ответ .

Answer 2

y, x = np.ogrid[:C.shape[0],:C.shape[1]]
D = A[C[y, x], y, x]