CODEWITHSUNDEEP
numpy
ovlindho
ovlindho

Reputation: 33

How to convert 4D numpy array to 2D by fixing indices

I have a 4D numpy array A of shape (N,N,N,N) that I would like to convert to a 2D matrix M of shape (N,N) by fixing pairs of indices. For example

M[i,j] = A[i,j,i,j]

How should this be done in numpy, avoiding for loops?

Edit: I will subsequently access the elements of M using an index array provided by numpy.ix_ so accessing elements of the 4D array in analoguous way would be a solution as well.

Upvotes: 2

Views: 171

Answers (1)

swag2198
swag2198

Reputation: 2696

This is a workaround:

i, j = np.arange(N), np.arange(N)
j_idx, i_idx = np.meshgrid(i, j)
M = A[i_idx, j_idx, i_idx, j_idx]

Uses meshgrid to generate the indexing pattern beforehand and then fancy indexing the array A to get M. As @hpaulj suggested, you can specify sparse = True in np.meshgrid() to obtain broadcastable 1D arrays instead of full 2D index arrays to save some space.

You can also do this using np.ix_() as well:

ixgrid = np.ix_(i, j)
M = A[ixgrid + ixgrid]

Since ixgrid is a 2-tuple, ixgrid + ixgrid produces the 4-tuple required for indexing A.

Upvotes: 2

Related Questions