CODEWITHSUNDEEP
pythonnumpyvectorization
Brandon Dube
Brandon Dube

Reputation: 448

Trace max along dim1 with varying index in dim3

I have a "cube" of 3D data where there is some peak in the column, or first dimension. The index of the peak may shift depending what row is examined. The third dimension may do something a bit more complicated, but for now can be thought of as just scaling things by some linear function.

I would like to find the index of the max along the first dimension, subject to the constraint that for each row, the z index is chosen such that the column peak will be closest to 0.5.

Here's a sample image that is a plane in row,column with a fixed z:

sample image of a row,column plane of the data with fixed z

These arrays will at times be large -- say, 21x11x200 float64s, so I would like to vectorize this calculation. Written with a for loop, it looks like this:

cols, rows, zs = data.shape
for i in range(rows):
    # for each field point, make an intermediate array that is 2D with focus,frequency dimensions
    arr = data[:,i,:]

    # compute the thru-focus max and find the peak closest to 0.5
    maxs = np.max(arr, axis=0)
    max_manip = np.abs(maxs-0.5)
    freq_idx = np.argmin(max_manip)

    # take the thru-focus slice that peaks closest to 0.5
    arr2 = data[:,i,freq_idx]
    focus_idx = np.argmax(arr2)
    print(focus_idx)

My issue is that I do not know how to roll these calculations up into a vector operation. I would appreciate any help, thanks!

Upvotes: 1

Views: 43

Answers (1)

Divakar
Divakar

Reputation: 221574

We just need to use the axis param with the relevant ufuncs there and that would lead us to a vectorized solution, like so -

# Get freq indices along all rows in one go
idx = np.abs(data.max(0)-0.5).argmin(1)

# Index into data with those and get the argmax indices
out = data[:,np.arange(data.shape[1]), idx].argmax(0)

Upvotes: 1

Related Questions