0

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?

smci
  • 32,567
  • 20
  • 113
  • 146
Vivek Sridhar
  • 37
  • 3

3 Answers3

2

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]])
xdze2
  • 3,986
  • 2
  • 12
  • 29
  • can this possibly be generalized to also work with the case where R = 1? say I have a 1x3x1 matrix, this command is giving me only the first element, not the full row. – Vivek Sridhar Aug 23 '18 at 15:16
  • 1
    It works if the matrix is 3x1x1... Which dimension corresponds to the layers? or you can try `np.diagonal(a, offset=0, axis1=0, axis2=2)` – xdze2 Aug 23 '18 at 15:22
  • That did the trick! (the layers are the third dimension). Thanks! – Vivek Sridhar Aug 23 '18 at 15:29
0

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]
J. Blackadar
  • 1,821
  • 1
  • 11
  • 18
0

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])]
MartyFizzle
  • 25
  • 6