沿轴累积量的最大值

2024-03-28 23:19:14 发布

您现在位置:Python中文网/ 问答频道 /正文

我有一个numpy数组aa.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]

必须有一个更干净、更快的方法来做到这一点。有什么想法吗?在


Tags: 代码infornp符号rangeabs数组
1条回答
网友
1楼 · 发布于 2024-03-28 23:19:14

我找不到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]

相关问题 更多 >