NumPy：如何使用 argmin/max 输出沿单个维度折叠 N 维数组？

Question

是否有一种直接的方法可以使用在 N 维数组的单个维度上调用 NumPy 的 argmax 或 argmin 函数的输出来定义该数组的索引？

这可能最好用一个例子来解释。考虑以下示例 随时间变化的二维温度读数网格：

>>> import numpy as np
>>> times = np.array([0, 20])
>>> temperature_map_t0 = np.array([[10, 12, 14], [23, 40, 50]])
>>> temperature_map_t1 = np.array([[20, 12, 15], [23, 10, 12]])
>>> temperature_map = np.dstack([temperature_map_t0, temperature_map_t1])

和包含相应压力读数的形状相同的 N-D 阵列：

>>> pressure_map = np.random.rand(*temperature_map.shape)

我们可以找到每个位置的最高温度：

>>> top_temperatures = temperature_map.max(axis=2)
>>> top_temperatures
array([[20, 12, 15],
       [23, 40, 50]])

以及它们发生的时间：

>>> times = times[temperature_map.argmax(axis=2)]
>>> times
array([[20,  0, 20],
       [ 0,  0,  0]])

但是我们如何使用temperature_map.argmax(axis=2)来找到对应的 压力？

>>> pressures_at_top_temperatures = pressures[ ???? ]

换句话说 - 折叠单个维度的索引语法是什么 使用该维度的 argmin 或 argmax 索引的 N 维数组？

Answer 1

感谢Jaime 当我有similar problem 时回答的问题

import numpy as np
times = np.array([0, 20])
temperature_map_t0 = np.array([[10, 12, 14], [23, 40, 50]])
temperature_map_t1 = np.array([[20, 12, 15], [23, 10, 12]])
temperature_map = np.dstack([temperature_map_t0, temperature_map_t1])
top_temperatures = temperature_map.max(axis=2)

# shape is a tuple - no need to convert
# http://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.shape.html
pressure_map = np.random.rand(*temperature_map.shape)

idx = temperature_map.argmax(axis=2)

s = temperature_map.shape
result pressure_map[np.arange(s[0])[:, None], np.arange(s[1]), idx]

Answer 2

我能想到的最直接的解决方案是使用逻辑索引将未由所需索引选择的条目归零，然后对感兴趣的维度求和，例如如下：

def collapse_dimension(ndarr, index, axis):
    r = np.rollaxis(ndarr, axis, 0)
    return np.sum((r[i] * (index == i) for i in range(r.shape[0])), axis=0)

所以对于上面的例子，我们可以使用argmax 或argmin 来折叠任意给定维度的数组，例如

>>> pressures_at_top_temperatures = collapse_dimension(
...     pressure_map, temperature_map.argmax(axis=2), 2)

并且，很简单，使用相应的argmax 在任何给定维度上获取max：

>>> temperature_map.max(axis=2) == collapse_dimension(
...     temperature_map, temperature_map.argmax(axis=2), 2)
array([[ True,  True,  True],
       [ True,  True,  True]], dtype=bool)

但是，我强烈怀疑有一种更好的方法可以做到这一点，而无需编写这个额外的函数——有什么想法吗？？