Numpy signed maximum magnitude of cumsum along an axis

Question

I have a numpy array a, a.shape=(17,90,144). I want to find the maximum magnitude of each column of cumsum(a, axis=0), but retaining the original sign. In other words, if for a given column a[:,j,i] the largest magnitude of cumsum corresponds to a negative value, I want to retain the minus sign.

The code np.amax(np.abs(a.cumsum(axis=0))) gets me the magnitude, but doesn't retain the sign. Using np.argmax instead will get me the indices I need, which I can then plug into the original cumsum array. But I can't find a good way to do so.

The following code works, but is dirty and really slow:

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]

There must be a cleaner, faster way to do this. Any ideas?

Answer 1

I can't find any alternatives to argmax but at least you can fasten that with a more vectorized approach:

# 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]