CODEWITHSUNDEEP
pythonarraysnumpymultidimensional-arrayindexing
fstab
fstab

Reputation: 5039

What is the best way to do multi-dimensional indexing with numpy?

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?

Upvotes: 1

Views: 89

Answers (1)

plonser
plonser

Reputation: 3363

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.

Upvotes: 2

Related Questions