Numpy 带符号的 cumsum 沿轴的最大幅度

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]