如果我有这样的 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(axis=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=...)
)提取索引,并使用这些索引从一堆其他(等效形状)数组中提取数据。我对实现此目的的替代方法持开放态度(例如,使用布尔掩码数组)。但是,我喜欢使用这些“索引”数组获得的“安全性”。有了这个,我保证有正确数量的元素来创建一个新数组,它看起来像一个通过 3d 字段的 2d“切片”。