NumPy:如何使用argmin/max输出沿单个维度折叠n维数组

Question

是否有一种直接的方法来使用在N-D数组的单个维度上调用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[ ???? ]

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

Answer 1

感谢Jaime，当我遇到类似的问题时，他回答了我

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)

然而，我有一个强烈的怀疑，有一个更好的方法来做到这一点，不涉及编写这个额外的函数-任何想法??