[Numpy-discussion] 2 greatest values, in a 3-d array, along one axis

[Angus McMorland]

On 3 August 2012 11:18, Jim Vickroy <jim.vickroy at noaa.gov> wrote:

> Hello everyone,
>
> I'm trying to determine the 2 greatest values, in a 3-d array, along one
> axis.
>
> Here is an approach:
>
> # ------------------------------------------------------
> # procedure to determine greatest 2 values for 3rd dimension of 3-d
> array ...
> import numpy, numpy.ma
> xcnt, ycnt, zcnt   = 2,3,4 # actual case is (1024, 1024, 8)
> p0                 = numpy.empty ((xcnt,ycnt,zcnt))
> for z in range (zcnt) : p0[:,:,z] = z*z
> zaxis              = 2                                            # max
> values to be determined for 3rd axis
> p0max              = numpy.max (p0, axis=zaxis)                   # max
> values for zaxis
> maxindices         = numpy.argmax (p0, axis=zaxis)                #
> indices of max values
> p1                 = p0.copy()                                    # work
> array to scan for 2nd highest values
> j, i               = numpy.meshgrid (numpy.arange (ycnt), numpy.arange
> (xcnt))
> p1[i,j,maxindices] = numpy.NaN                                    # flag
> all max values
> p1                 = numpy.ma.masked_where (numpy.isnan (p1), p1) # hide
> all max values
> p1max              = numpy.max (p1, axis=zaxis)                   # 2nd
> highest values for zaxis
> # additional code to analyze p0max and p1max goes here
> # ------------------------------------------------------
>
> I would appreciate feedback on a simpler approach -- e.g., one that does
> not require masked arrays and or use of magic values like NaN.
>
> Thanks,
> -- jv
>

Here's a way that only uses argsort and fancy indexing:

>>>a = np.random.randint(10, size=(3,3,3))
>>>print a

[[[0 3 8]
  [4 2 8]
  [8 6 3]]

 [[0 6 7]
  [0 3 9]
  [0 9 1]]

 [[7 9 7]
  [5 2 9]
  [9 3 3]]]

>>>am = a.argsort(axis=2)
>>>maxs = a[np.arange(a.shape[0])[:,None], np.arange(a.shape[1])[None],
am[:,:,-1]]
>>>print maxs

[[8 8 8]
 [7 9 9]
 [9 9 9]]

>>>seconds = a[np.arange(a.shape[0])[:,None], np.arange(a.shape[1])[None],
am[:,:,-2]]
>>>print seconds

[[3 4 6]
 [6 3 1]
 [7 5 3]]

And to double check:

>>>i, j = 0, 1
>>>l = a[i, j,:]
>>>print l

[4 2 8]

>>>print np.max(a[i,j,:]), maxs[i,j]

8 8

>>>print l[np.argsort(l)][-2], second[i,j]

4 4

Good luck.

Angus.
-- 
AJC McMorland
Post-doctoral research fellow
Neurobiology, University of Pittsburgh
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20120803/fe42f4aa/attachment.html>