For loop in numpy 3d arrays

Question

I want to take a basic 3d array like this:

b = np.arange(1,101).reshape(4,5,5)
b

Then I want to take the first index, and work down like a stairs.

b1 = [b[0:,0,0],b[0:,1,1],b[0:,2,2],b[0:,3,3],b[0:,4,4]]
b1 = np.asarray(b1)
b1 = np.transpose(b1)
b1

The above code doesn't look right, I'd rather use a loop. This is what I have so far:

for i in range(0,5):
    b2 = b[0:,i,i]
    b2 = np.asarray(b2)
    b2 = b2.reshape(4,1)
    print(b2)

My issue with the above output is it puts each iteration into one vertical array, then moves onto the next. How do I make the above code output something similar to my second block of code?

Apologies for the poor formatting, new to stackoverflow and just starting to learn code/numpy.

Thanks!

Answer 1

You may use list comprehension:

b1 = [b[0:,i,i] for i in range(5)]
b1 = np.asarray(b1)
b1 = np.transpose(b1)

Answer 2

You can use einsum to extract all the stairs at once:

>>> np.einsum("ijj->ij",b) 
array([[  1,   7,  13,  19,  25],
       [ 26,  32,  38,  44,  50],
       [ 51,  57,  63,  69,  75],
       [ 76,  82,  88,  94, 100]])

and then split into columns:

>>> np.split(np.einsum("ijj->ij",b),np.arange(1,5),1)
[array([[ 1],
       [26],
       [51],
       [76]]), array([[ 7],
       [32],
       [57],
       [82]]), array([[13],
       [38],
       [63],
       [88]]), array([[19],
       [44],
       [69],
       [94]]), array([[ 25],
       [ 50],
       [ 75],
       [100]])]

Answer 3

I think this is what you want to do

import numpy as np

b = np.arange(1,101).reshape(4,5,5)

result = np.diag(b[0])

print(result)

result:

[ 1  7 13 19 25]

Answer 4

Another way of doing it (presumably in your stair, b has to have square shape in its last two dimensions):

c = b[:,np.arange(b.shape[1]),np.arange(b.shape[2])]
b2 = c.T.reshape(c.shape+(1,))

output:

[[[  1]
  [ 26]
  [ 51]
  [ 76]
  [  7]]

 [[ 32]
  [ 57]
  [ 82]
  [ 13]
  [ 38]]

 [[ 63]
  [ 88]
  [ 19]
  [ 44]
  [ 69]]

 [[ 94]
  [ 25]
  [ 50]
  [ 75]
  [100]]]