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.
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.