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.
import numpy as np
arr = np.array([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]])
print(np.diagonal(arr, axis1=1, axis2=2))
Output
[[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.
import numpy as np
# Create a 3D numpy array of shape (3, 4, 4)
arr = np.arange(3 * 4 * 4).reshape(3, 4, 4)
print("Original 3D array:\n", arr)
# Extract diagonals from each 2D sub-array
diag_arr = np.diagonal(arr, axis1=1, axis2=2)
print("\n2D diagonal array:\n", diag_arr)
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.
import numpy as np
# Create a 3D numpy array of shape (3, 3, 4)
arr = np.arange(3 * 3 * 4).reshape(3, 3, 4)
print("Original 3D array:\n", arr)
# Extract diagonals
diag_arr = np.diagonal(arr, axis1=1, axis2=2)
print("\n2D diagonal array:\n", diag_arr)
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.
import numpy as np
# Create a 3D numpy array of shape (3, 5, 6)
arr = np.arange(3 * 5 * 6).reshape(3, 5, 6)
print("Original 3D array:\n", arr)
# Extract diagonals
diag_arr = np.diagonal(arr, axis1=1, axis2=2)
print("\n2D diagonal array:\n", diag_arr)
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).