python - 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维数组:

>>> 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]