从一个数组计算 argmax 并用于从另一个数组中获取值

Question

我正在尝试从一个 ndarray 获取 argmax 并使用它从另一个 ndarray 获取值，但我做错了。

ndvi_array = np.random.randint(0, 255, size=(4, 1, 100, 100))
image_array = np.random.randint(0, 255, size=(4, 12, 100, 100))
ndvi_argmax = ndvi_array.argmax(0)
print(f"NDVI argmax shape: {ndvi_argmax.shape}")
zipped = tuple(zip(range(len(ndvi_argmax)), ndvi_argmax))
result = image_array[zipped]
print(f"Result share: {result.shape}")

我收到以下错误：

only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

如何获得具有最大值的形状数组 (1,12,100,100)？

Answer 1

>>> result = np.take_along_axis(
...     image_array,
...     ndvi_array.argmax(axis=0, keepdims=True),
...     axis=0,
    )

>>> print(f"{result.shape = }")
result.shape = (1, 12, 100, 100)