我有四个矩阵，部分由彼此创建：
A是表示灰度图像堆栈的3D矩阵，其形状为(n, h, w)B是也表示图像堆栈的3D矩阵，其中每个切片是根据A中的相应切片单独计算的，并且也是形状为(n, h, w)C是2D矩阵，包含z方向上最大值为B的索引，形状为(h, w)D是2D矩阵，其中来自a的值从某个切片复制，该切片由位置(x, y)处的C中的值指示
用循环实现的最小示例如下：
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]])
我的问题是：如何用C对A进行切片，有效地消除嵌套的for循环？我最初的想法是将C扩展到一个3D布尔掩码，其中只有位置[c, y, x]设置为True，然后简单地将其与A元素相乘，并在z轴上求和。但我想不出一个没有循环的Python式实现(如果我知道如何编写的话，我可能就不需要再创建布尔掩码了(。
我发现最接近已经实现的函数是np.chose((，但它只需要32个元素就可以实现C.
这里的标准方法是将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
表明这是迄今为止提出的最有效的方法。还要注意，与_fix方法相反，基于np.take_along_axis()和take_by_axis_ogrid()方法对于具有更高维度的输入基本上不变。
特别地，take_by_axis_ogrid()是take_by_axis_ogrid_fix()的axis不可知版本，本质上是第n个答案。
y, x = np.ogrid[:C.shape[0],:C.shape[1]]
D = A[C[y, x], y, x]