Question

我有一个numpy数组a，a.shape=(17,90,144)。我想找到cumsum(a, axis=0)每列的最大幅度，但保留原始符号。换句话说，如果对于给定列a[:,j,i]，cumsum的最大幅度对应于负值，我想保留减号。

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

Numpy沿轴线签署了最大量的cumsum

1 个答案: