CODEWITHSUNDEEP
pythonnumpy
gian9
gian9

Reputation: 107

Indexing multidimensional arrays with argmin indices applied along one axis

I have two multidimensional array, let's say x and y both of them with 5 dimensions and I would like to find the value of x at which the last component of y is minima. To find the indexes I just use I=argmin(y,axis=-1) and this returns me a 4-dimensional array of indexes. How should I do to find the values of x for those indexes? Some sort of x[I]?

Upvotes: 3

Views: 388

Answers (2)

hpaulj
hpaulj

Reputation: 231335

Using argmin with a 1 or 2d array is fairly simple, but with 3 or more, the mapping is harder to understand:

In [332]: y=np.arange(24)
In [333]: np.random.shuffle(y)
In [334]: y=y.reshape(2,3,4)
In [335]: y
Out[335]: 
array([[[19, 12,  9, 21],
        [ 8, 13, 20, 17],
        [22, 11,  5,  1]],

       [[ 7,  2, 23, 16],
        [ 0, 10,  6,  4],
        [14, 18, 15,  3]]])

In [338]: I = np.argmin(y, axis=-1)
In [339]: I
Out[339]: 
array([[2, 0, 3],
       [1, 0, 3]], dtype=int32)
In [340]: np.min(y, axis=-1)
Out[340]: 
array([[9, 8, 1],
       [2, 0, 3]])

The result is (2,3), one index for each plane/row.

The I[0,0] means that y[i,j,I[i,j]] is the minimum in the i,j row.

So we need a way to generate that i,j pairing

In [345]: i,j = np.ix_(np.arange(2), np.arange(3))
In [346]: i
Out[346]: 
array([[0],
       [1]])
In [347]: j
Out[347]: array([[0, 1, 2]])

In [349]: y[i,j,I[i,j]]
Out[349]: 
array([[9, 8, 1],
       [2, 0, 3]])

Or shorten that to:

In [350]: y[i,j,I]
Out[350]: 
array([[9, 8, 1],
       [2, 0, 3]])

Even with 2d the method is the same:

In [360]: z=y[:,:,1]
In [361]: z
Out[361]: 
array([[12, 13, 11],
       [ 2, 10, 18]])
In [362]: idx=np.argmin(z, axis=-1)
In [363]: idx
Out[363]: array([2, 0], dtype=int32)
In [364]: z[[0,1], idx]       # index the 1st dim with range
Out[364]: array([11,  2])

Using mgrid might make it easier to visualize the process:

In [378]: i,j =np.mgrid[0:2,0:3]
In [379]: i
Out[379]: 
array([[0, 0, 0],
       [1, 1, 1]])
In [380]: j
Out[380]: 
array([[0, 1, 2],
       [0, 1, 2]])
In [381]: y[i, j, I]
Out[381]: 
array([[9, 8, 1],
       [2, 0, 3]])

here i and j are (2,3) arrays that match I in shape. Together the 3 arrays select a (2,3) array of elements from y.

ix_ and ogrid just generate the equivalent open arrays.

Upvotes: 2

Divakar
Divakar

Reputation: 221504

Approach #1 : That's basically advanced-indexing extended to 5D case. To make things a bit more convenient, we can make use of open range arrays with np.ogrid and then perform the advanced-indexing, like so -

d0,d1,d2,d3,d4 = x.shape
s0,s1,s2,s3 = np.ogrid[:d0,:d1,:d2,:d3]
ymin = y[s0,s1,s2,s3,I]
xmin = x[s0,s1,s2,s3,I]

Approach #2 : We can shorten it a bit, by merging the first two steps with np.ix_ and thus have a generic function to handle ndarrays of generic number of dimensions -

indxs = np.ix_(*[np.arange(i) for i in x.shape[:-1]]) + (I,)
ymin = y[indxs]
xmin = x[indxs]

Let's use some sample random valued array and verify by directly computing the min along last axis with y.min(axis=-1) i.e. y.min(-1) and comparing it against the indexed value ymin from the proposed codes -

In [117]: x = np.random.randint(0,9,(3,4,5,6,7))
     ...: y = np.random.randint(0,9,(3,4,5,6,7))
     ...: I = np.argmin(y,axis=-1)
     ...: 

In [118]: d0,d1,d2,d3,d4 = x.shape
     ...: s0,s1,s2,s3 = np.ogrid[:d0,:d1,:d2,:d3]
     ...: ymin = y[s0,s1,s2,s3,I]
     ...: xmin = x[s0,s1,s2,s3,I]
     ...: 

In [119]: np.allclose( y.min(-1), ymin)
Out[119]: True

In [120]: indxs = np.ix_(*[np.arange(i) for i in x.shape[:-1]]) + (I,)
     ...: ymin = y[indxs]
     ...: xmin = x[indxs]
     ...: 

In [121]: np.allclose( y.min(-1), ymin)
Out[121]: True

Upvotes: 3

Related Questions