Numpy 沿某轴的 cumsum 的有符号最大幅度

Question

我有一个numpy数组 a，它的形状是 a.shape=(17,90,144)。我想找到每一列在 cumsum(a, axis=0) 中的最大绝对值，但要保留原来的符号。换句话说，如果某一列 a[:,j,i] 的最大绝对值对应的是一个负数，我希望保留这个负号。

代码 np.amax(np.abs(a.cumsum(axis=0))) 可以让我得到绝对值，但它不保留符号。使用 np.argmax 可以让我找到需要的索引，然后我可以把这些索引放回原来的 cumsum 数组中。但是我找不到一个好的方法来做到这一点。

下面的代码可以实现这个功能，但写得很乱，而且速度非常慢：

max_mag_signed = np.zeros((90,144))
indices = np.argmax(np.abs(a.cumsum(axis=0)), axis=0)
for j in range(90):
    for i in range(144):
        max_mag_signed[j,i] = a.cumsum(axis=0)[indices[j,i],j,i]

一定有更简洁、更快速的方法来实现这个功能。有没有好的建议？

Answer 1

我找不到argmax的其他替代方法，但至少你可以用更快的向量化方法来加速这个过程：

# store the cumsum, since it's used multiple times
cum_a = a.cumsum(axis=0)

# find the indices as before
indices = np.argmax(abs(cum_a), axis=0)

# construct the indices for the second and third dimensions
y, z = np.indices(indices.shape)

# get the values with np indexing
max_mag_signed = cum_a[indices, y, z]