在 3D 数组的特定轴上查找轴上每组值的最大值

Question

我有一个形状为 (8, x, y) 的 3D numpy 数组，其中 x 和 y 可以是 0 到 ~500 之间的任何整数值。这 8 个形状为 (x, y) 的“层”各自表示平面上特定点处可能的 z 值（相应的坐标单独存储）。 z 值可以为正值或负值。

如何找到每个坐标点处的signed绝对最大值？最好不要循环 - 这些层可能会变得相当大。

简单的例子：

>>> a = np.array([
    [
      [3, 5],
      [-4, 1],
    ],
    [
      [1, 2],
      [1, 4],
    ],
    [
      [2, -3],
      [0, -5],
    ],
    [
      [-6, 4],
      [2, 2],
    ]
])

预期结果：

[[-6, 5],
 [-4, -5]]

这是我当前的工作示例：

a = np.array([
    [
        [3, 5],
        [-4, 1],
    ],
    [
        [1, 2],
        [1, 4],
    ],
    [
        [2, -3],
        [0, -5],
    ],
    [
        [-6, 4],
        [2, 2],
    ],
])

a_max = np.max(a, axis=0)
a_min = np.min(a, axis=0)

# Stack max and min for later masking via argmax
a_max_min = np.dstack([a_max, a_min])

# find max and min value indices
a_max_min_argmax = np.argmax(a_max_min, axis=0)
a_max_min_argmin = np.argmin(a_max_min, axis=0)

# Find absolute max array using mask
a_abs_max = a_max * a_max_min_argmax + a_min * a_max_min_argmin

结果：

[[-6  5]
 [-4 -5]]

这可行，但非常粗糙。我也知道如果有例如它会破裂。 -5 和 5 在同一点。

我该如何改进？我认为应该有一种方法可以使用

np.argmax(np.abs(a_max_min), axis=0)

，但我不知道如何去做。

提前致谢！

Answer 1

看来您对 argmax 方法的看法绝对正确。 你的代码非常接近

这是最终代码：

indices = np.argmax(np.abs(a), axis=0)
max_z = np.take_along_axis(a, indices[np.newaxis, ...], axis=0)

print(indices)
print(max_z)

argmax 返回最大值所在的索引。 take_along_axis 只是循环并根据索引获取所需的数字。

如果打印它，您可以看到它是如何工作的（例如，索引的左上角不是显示 -6，而是显示 3。-6 是我们想要的数字，位于第三个二维数组中） .

说实话，我不确定避免 for 循环是否是解决这些问题的正确方法？我觉得 Python 处理这些东西已经足够高效了，除非我们谈论的是数十亿个数据点。

但是学习使用 numpy 函数来做到这一点仍然很酷。