Search code examples
pythonnumpymatrixmatrix-indexing

Selectively access rows of numpy 3D matrix

Is there a way to index a 3D numpy matrix to selectively grab the i'th row of each i'th layer?

e.g. I have an RxNxR matrix, and I want to grab the 1st row of the 1st layer, the 2nd row of the 2nd layer, the 3rd row of the 3rd layer, etc. and end up with an RxN matrix.

Can I do that in a single operation?

Solution

  • Using the numpy.diagonal function:

    a = np.arange(27).reshape(3, 3, 3)
np.diagonal(a, offset=0, axis1=0, axis2=1).T

    gives

    array([[ 0,  1,  2],
       [12, 13, 14],
       [24, 25, 26]])