用numpy进行多维索引的最佳方法是什么？

Question

I am trying to do some indexing on a 3D numpy array. Basically I have an array phi which has shape (F,A,D) ; for example (5, 3, 7) . Generated, for example as follows:

F=5; A=3; D=7; phi = np.random.random((F,A,D))

My goal is to be able to index over A and D , with a 2D array such as [[0,1,2],[5,5,6]] , which means take the values indexed by 0 in the 3rd dimension, for the the first position in A , the values indexed by 1 in the 3rd dimension for the second position of A and so on. The result should have a shape that is (F,A,2) or (F,2,A) .

This would be equivalent to manually cycling all the values of the "indexer array" such as:

phi[:,0,0]; phi[:,1,1]; phi[:,2,2]
phi[:,0,5]; phi[:,1,5]; phi[:,2,6]

Intuitively I would do something like phi[:,:,[[0,1,2],[3,3,3]]] , but it's shape ends up being (5, 3, 2, 3) .

Any ideas on how to obtain the correct result?

Answer 1

I think this is what you want

phi[:,range(A),[[0,1,2],[5,5,6]]]

Your attempt

phi[:,:,[[0,1,2],[5,5,6]]]

takes the values along the third dimension for every values of the first two dimensions, therefore you end up with a shape of (5,3,2,3) .

However, according to your example you want a continous increase in the second dimension which is accomplished in my code by range(A) and numpy's broadcasting.