Qfyilyi

So I have created an array from 3 nested lists (atleast I think it is an array from 3 lists), and I want to access the three diagonal elements in it. I have the array created, but how do I access the three diagonal elements in it?

from numpy import *

test1 = arange(27).reshape(3,3,3)

test1

Result:

array([[[ 0,  1,  2],

      [ 3,  4,  5],

      [ 6,  7,  8]],



     [[ 9, 10, 11],

      [12, 13, 14],

      [15, 16, 17]],



     [[18, 19, 20],

      [21, 22, 23],

      [24, 25, 26]]])

Here is a list comprehension approach:

>>> [np.diagonal(i) for i in test1]

[array([0, 4, 8]), array([ 9, 13, 17]), array([18, 22, 26])]

There are several ways to achieve your goal. Here, I will highlight the use of a boolean mask.

First create the boolean 3x3 identity matrix : i.e. the diagonal is True whilst 2. every off diagonal entry is False.
Then overlay the boolean mask over your original ndarray to get the diagonals.

import numpy as np

test1 = np.arange(27).reshape(3,3,3)



>>> diag = np.eye(3, dtype=bool)

>>> test1[:, diag]

array([[ 0,  4,  8],

       [ 9, 13, 17],

       [18, 22, 26]])

As you can see, this gives a 2d array where each row is the corresponding diagonal of the zeroth, first and second
2d array in your 3d array.

As an aside, avoid import *, it is the cause of many a headache because if destroys the namespace abstraction you
have. In the above example, what if numpy had a diag function or variable defined? same if you import another package after numpy and it happens to have it's own arange function, you will looes numpy's arange function.
Prefer explicit imports to star imports.