How to get all 2D diagonals of a 3D NumPy array

In NumPy, you can extract 2D diagonals from each sub-array of a 3D array using numpy.diagonal() function.

Let’s start with a very simple example to see how it works.

[[1 4]
 [5 8]]

Explanation:

• From the first 2×2 block -> diagonal is [1, 4].
• From the second block -> diagonal is [5, 8].

So the result is a 2D array containing diagonals from each block.

diagonal Function

This function allows you to extract diagonals from an n-dimensional array.

Syntax

np.diagonal(a, axis1, axis2)

Parameters:

• a: The array from which diagonals will be taken.
• axis1: First axis that defines the 2D sub-arrays.
• axis2: Second axis that defines the 2D sub-arrays.

Return: A new array containing the diagonals of the selected sub-arrays.

Examples

Example 1: In this example, we create a 3D array with shape (3,4,4). Each "block" is a 4×4 2D matrix and we’ll grab the diagonals from each one.

Output

Original 3D 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 27]
[28 29 30 31]]

[[32 33 34 35]
[36 37 38 39]
[40 41 42 43]
[44 45 46 47]]]

2D diagonal array:
[[ 0 5 10 15]
[16 21 26 31]
[32 37 42 47]]

Explanation: Here, NumPy takes the diagonals of each 4×4 block:

• From the 1st block: [0, 5, 10, 15]
• From the 2nd block: [16, 21, 26, 31]
• From the 3rd block: [32, 37, 42, 47]

So the final result is a 2D array of diagonals.

Example 2: In this example, we create a 3D array with shape (3,3,4). Each block is 3×4 and NumPy extracts diagonals across them.

Output

Original 3D 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 27]
[28 29 30 31]
[32 33 34 35]]]

2D diagonal array:
[[ 0 5 10]
[12 17 22]
[24 29 34]]

Explanation:

• From the 1st 3×4 block: diagonal [0, 5, 10]
• From the 2nd block: [12, 17, 22]
• From the 3rd block: [24, 29, 34]

The diagonals are stacked into a 2D array of shape (3,3).

Example 3: In this example, we use a bigger array of shape (3,5,6). Each block is 5×6 and the diagonals are extracted across all blocks.

Output

Original 3D 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 27 28 29]]

[[30 31 32 33 34 35]
[36 37 38 39 40 41]
[42 43 44 45 46 47]
[48 49 50 51 52 53]
[54 55 56 57 58 59]]

[[60 61 62 63 64 65]
[66 67 68 69 70 71]
[72 73 74 75 76 77]
[78 79 80 81 82 83]
[84 85 86 87 88 89]]]

2D diagonal array:
[[ 0 7 14 21 28]
[30 37 44 51 58]
[60 67 74 81 88]]

Explanation:

• From the 1st 5×6 block: diagonal [0, 7, 14, 21, 28]
• From the 2nd block: [30, 37, 44, 51, 58]
• From the 3rd block: [60, 67, 74, 81, 88]

So NumPy gives us a neat 2D diagonal array with shape (3,5).