Selective Grab Rows of Matrix Numpy

Is there a way to index a 3D numpy matrix to selectively grab the ith row of each 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?

3 Answers
3

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]])


np.diagonal(a, offset=0, axis1=0, axis2=2)

I believe this should do the trick:

import numpy as np
matrix = ...
result = np.zeros((matrix.shape[0], matrix.shape[1]))

for i in range(0, len(matrix)):
result[i] = matrix[i][i]


I think you could use a comprehension - depends on which dimensions you define as your 'layer', etc but this would work:

a = np.array([...])

[a[:, i, 0] for i in range(np.shape(a)[1])]


By clicking "Post Your Answer", you acknowledge that you have read our updated terms of service, privacy policy and cookie policy, and that your continued use of the website is subject to these policies.