CODEWITHSUNDEEP
numpymaxargmax
L. Marval
L. Marval

Reputation: 25

In numpy, what is the efficient way to find the maximum values and their indices of a 3D ndarray across two axis?

How to find the correlation-peak values and coordinates of a set of 2D cross-correlation functions? Given an 3D ndarray that contains a set of 2D cross-correlation functions. What is the efficient way to find the maximum(peak) values and their coordinates(x and y indices)?
The code below do the work but I think it is inefficient. 

import numpy as np
import numpy.matlib 


ccorr  = np.random.rand(7,5,5)
xind   = ccorr.argmax(axis=-1)
mccorr = ccorr[np.matlib.repmat(np.arange(0,7)[:,np.newaxis],1,5),np.matlib.repmat(np.arange(0,5)[np.newaxis,:],7,1), xind]
yind   = mccorr.argmax(axis=-1)
xind   = xind[np.arange(0,7),yind]
values = mccorr[np.arange(0,7),yind]

print("cross-correlation functions (z,y,x)")
print(ccorr)

print("x and y indices of the maximum values")
print(xind,yind)
print("Maximum values")
print(values)

Upvotes: 2

Views: 87

Answers (1)

Daniel F
Daniel F

Reputation: 14399

You'll want to flatten the dimensions you're searching over and then use unravel_index and take_along_axis to get the coordinates and values, respectively.

ccorr  = np.random.rand(7,5,5)
cc_rav = ccorr.reshape(ccorr.shape[0], -1)
idx = np.argmax(cc_rav, axis = -1)
indices_2d = np.unravel_index(idx, ccorr.shape[1:])
vals = np.take_along_axis(ccorr, indices = indices_2d, axis = 0)

if you're using numpy version <1.15:

vals = cc_rav[np.arange(ccorr.shape[0]), idx]

or:

vals = ccorr[np.arange(ccorr.shape[0]), 
             indices_2d[0], indices_2d[1]]

Upvotes: 1

Related Questions