CODEWITHSUNDEEP
pythonarraysnumpyindexing
Eric
Eric

Reputation: 97571

Open mesh of array indices

Given a matrix a where a.shape == (M, N, O), does there exist a better way to generate:

indices = (
    np.arange(M).reshape(M, 1, 1),
    np.arange(N).reshape(1, N, 1),
    np.arange(O).reshape(1, 1, O)
)

And also works for higher-dimensions of array?

I can get close with np.indices, but this returns a closed mesh (all entries are of shape M,N,O), not an open mesh.

Upvotes: 3

Views: 351

Answers (2)

ali_m
ali_m

Reputation: 74172

np.ogrid is probably the most concise method if M, N and O are specified separately:

M, N, O = 2, 3, 4

indices = np.ogrid[:M, :N, :O]

print(indices)
# [array([[[0]],

#        [[1]]]), array([[[0],
#         [1],
#         [2]]]), array([[[0, 1, 2, 3]]])]

If the input is just the array itself, the most concise way I can think of is:

np.ix_(*(np.r_[:s] for s in A.shape))

Upvotes: 2

Alex Riley
Alex Riley

Reputation: 176810

You can possibly do this with np.ix_:

np.ix_(np.arange(M), np.arange(N), np.arange(O))

From the documentation:

Construct an open mesh from multiple sequences.

This function takes N 1-D sequences and returns N outputs with N dimensions each, such that the shape is 1 in all but one dimension and the dimension with the non-unit shape value cycles through all N dimensions.

This produces the same result as indices here (M, N, O = 2, 3, 4):

(array([[[0]],
 
        [[1]]]), 

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

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

Upvotes: 3

Related Questions