Numpy index columns-CodePudding

I have a 2d matrix D

array([[1, 1, 8, 8],
       [3, 3, 4, 4],
       [8, 8, 1, 1]])

I would like to find each column's lowest value by indexing it with the amin array:

amin = D.argmin(0)
amin

array([0, 0, 2, 2])

How do I properly index D with amin such that the returned array is array([1, 1, 1, 1])?

CodePudding user response：

You need a combination of expand_dims and take_along_axis

np.take_along_axis(D, np.expand_dims(amin, 0), 0)
Out[]: array([[1, 1, 1, 1]])

This is arguably not the most intuitive formulation, but without going into design philosophy of numpy, I can say it's as "pythonic" as possible in this case.

CodePudding user response：

In [99]: D=np.array([[1, 1, 8, 8],
    ...:        [3, 3, 4, 4],
    ...:        [8, 8, 1, 1]])
In [100]: amin = D.argmin(0)
In [101]: amin
Out[101]: array([0, 0, 2, 2])

What is amin? It has one value per column, but that value is a row index. So what you want is:

In [102]: D[amin[0], 0]
Out[102]: 1
In [103]: D[amin[2], 2]
Out[103]: 1

Or doing all columns at once:

In [104]: D[amin, [0,1,2,3]]
Out[104]: array([1, 1, 1, 1])

Expressions like D[:, amin] try to apply amin to indexing the columns, not the rows. And the slice selects all rows at once. It doesn't pair them up.

May be these versions will make the action clearer:

In [105]: D[[0,0,2,2], [0,1,2,3]]
Out[105]: array([1, 1, 1, 1])
In [106]: [D[i,j] for i,j in zip([0,0,2,2],[0,1,2,3])]
Out[106]: [1, 1, 1, 1]

take_along_axis does the same thing as [104]

In [107]: np.take_along_axis(D, amin[None,:],0)
Out[107]: array([[1, 1, 1, 1]])

It's a bit cleaner if we tell argmin to keepdims (a new feature in v 1.22)

In [108]: amin = D.argmin(0, keepdims=True)
In [109]: amin
Out[109]: array([[0, 0, 2, 2]])
In [110]: np.take_along_axis(D, amin,0)
Out[110]: array([[1, 1, 1, 1]])