Indexing multidimensional arrays with argmin indices applied along one axis
I have two multidimensional arrays, let x and y both of them with 5 dimensions, and I would like to find the value of x in which the last component of y is the minimum. To find the indices I just use I=argmin(y,axis=-1)
and this returns me a 4d array of indices. How do I go about finding the x values for these indices? Some kind x[I]
?
source to share
Approach # 1: Mainly advanced-indexing
applies to 5D
. To make things more convenient, we can use open range arrays with np.ogrid
and then execute advanced-indexing
like -
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 using the first two steps with np.ix_
and therefore have a common function to handle ndarrays of the total 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 array and check by directly calculating min
along the last axis using y.min(axis=-1)
ie y.min(-1)
and comparing it to the indexed value ymin
with the suggested 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
source to share
Using a argmin
1 or 2d array is pretty straightforward, 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.
I[0,0]
means that y[i,j,I[i,j]]
is the minimum in the line i,j
.
So we need a way to generate pairing i,j
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 this 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])
Usage mgrid
can facilitate visualization of 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
- (2,3) arrays corresponding to the shape i
. Together 3 arrays select an array (2,3) of elements from y
.
ix_
and ogrid
just generate equivalent arrays open
.
source to share