python - Numpy:索引 3D 数组，最后一个轴的索引存储在 2D 数组中

Question

我有一个包含值的 shape(z,y,x) 的 ndarray。我试图用 shape(y,x) 的另一个 ndarray 索引这个数组，它包含我感兴趣的值的 z-index。

import numpy as np
val_arr = np.arange(27).reshape(3,3,3)
z_indices = np.array([[1,0,2],
                      [0,0,1],
                      [2,0,1]])

因为我的数组相当大，所以我尝试使用 np.take 来避免不必要的数组副本，但就是无法用它来索引 3 维数组。

如何使用 z_indices 索引 val_arr 以获得所需 z 轴位置的值？预期结果将是:

result_arr = np.array([[9,1,20],
                       [3,4,14],
                       [24,7,17]])

Answer 1

您可以使用 choose进行选择:

>>> z_indices.choose(val_arr)
array([[ 9,  1, 20],
       [ 3,  4, 14],
       [24,  7, 17]])

choose 函数非常有用，但理解起来可能有些棘手。本质上，给定一个数组 (val_arr)，我们可以从沿第一轴的每个 n 维切片做出一系列选择 (z_indices)。

此外:任何奇特的索引操作都会创建一个新数组而不是原始数据的 View 。在不创建全新数组的情况下，无法使用 z_indices 对 val_arr 进行索引。