python - 基于 2d 掩码数组的 numpy 3d 到 2d 转换

Question

如果我有这样的 ndarray:

>>> a = np.arange(27).reshape(3,3,3)
>>> a
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]],

       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]]])

我知道我可以使用 np.max(axis=...) 沿特定轴获得最大值:

>>> a.max(axis=2)
array([[ 2,  5,  8],
       [11, 14, 17],
       [20, 23, 26]])

或者，我可以沿该轴获取与最大值相对应的指数:

>>> indices = a.argmax(axis=2)
>>> indices
array([[2, 2, 2],
       [2, 2, 2],
       [2, 2, 2]])

我的问题——给定数组 indices 和数组 a，是否有一种优雅的方法来重现 a.max(轴=2)?

这可能会起作用:

import itertools as it
import numpy as np
def apply_mask(field,indices):
    data = np.empty(indices.shape)

    #It seems highly likely that there is a more numpy-approved way to do this.
    idx = [range(i) for i in indices.shape]
    for idx_tup,zidx in zip(it.product(*idx),indices.flat):
        data[idx_tup] = field[idx_tup+(zidx,)]
    return data

但是，它看起来很笨拙/效率低下。它也不允许我将它用于“最后”轴以外的任何轴。是否有一个 numpy 函数(或一些神奇的 numpy 索引的使用)来完成这项工作？天真的 a[:,:,a.argmax(axis=2)] 不起作用。

更新:

似乎以下也有效(并且更好一点):

import numpy as np
def apply_mask(field,indices):
    data = np.empty(indices.shape)

    for idx_tup,zidx in np.ndenumerate(indices):
        data[idx_tup] = field[idx_tup+(zidx,)]

    return data

---

我想这样做是因为我想根据 1 个数组中的数据提取索引(通常使用 argmax(axis=...))并使用这些索引来提取数据从一堆其他(等效形状)阵列中。我愿意接受其他方法来实现这一点(例如使用 bool 掩码数组)。但是，我喜欢使用这些“索引”数组所带来的“安全性”。有了这个，我保证有正确数量的元素来创建一个新数组，它看起来像 3d 字段中的 2d“切片”。

Answer 1

这是一些神奇的 numpy 索引，可以做你想做的事，但不幸的是它很难读。

def apply_mask(a, indices, axis):
    magic_index = [np.arange(i) for i in indices.shape]
    magic_index = np.ix_(*magic_index)
    magic_index = magic_index[:axis] + (indices,) + magic_index[axis:]
    return a[magic_index]

或者同样不可读:

def apply_mask(a, indices, axis):
    magic_index = np.ogrid[tuple(slice(i) for i in indices.shape)]
    magic_index.insert(axis, indices)
    return a[magic_index]