The diag()
method either creates a new ndarray
with the given 1D array as its diagonal elements or it extracts the diagonal from the given ndarray
.
Example
import numpy as np
# create a 1D array
array1 = np.arange(3)
# create a 2D array
array2 = np.arange(9).reshape(3,3)
# create a 2D array with elements of arr1 as diagonal
diagonalArray1 = np.diag(array1)
print('Array1:\n', array1)
print('Array1 as diagonal elements:\n', diagonalArray1)
# extract diagonal elements from arr2
diagonalArray2 = np.diag(array2)
print('\nArray2\n', array2)
print('Extract diagonal of Array2\n', diagonalArray2)
'''
Array1:
[0 1 2]
Array1 as diagonal elements:
[[0 0 0]
[0 1 0]
[0 0 2]]
Array2
[[0 1 2]
[3 4 5]
[6 7 8]]
Extract diagonal of Array2
[0 4 8]
'''
diag() Syntax
The syntax of diag()
is:
numpy.diag(array, k = 0)
diag() Arguments
The diag()
method takes the following arguments:
array
- input array (can bearray_like
)k
(optional) - an integer number representing the diagonal to retrieve
Note:
- By default
k
= 0 and represents the main diagonal. k
> 0 represents diagonals above the main diagonalk
< 0 represents diagonals below the main diagonal.
diag() Return Value
The diag()
method either returns a new ndarray with values on the 1D array as its diagonal, or returns a 1D array containing the diagonal elements of a given ndarray.
Example 1: Create a Diagonal Array With 1D Array
When a 1D array is passed to diag()
, it creates a diagonal array with the given array as diagonal elements.
As discussed earlier, we can use the k
argument to control the placement of the diagonal elements in the resulting array.
Let us look at an example.
import numpy as np
# create a 1D array
array1 = np.arange(1, 4)
# create a 2D array with elements of array1 as the main diagonal
mainDiagonal = np.diag(array1)
# create a 2D array with elements of array1 as diagonal above the main diagonal
upperDiagonal = np.diag(array1, k = 1)
# create a 2D array with elements of array1 as diagonal below the main diagonal
lowerDiagonal = np.diag(array1, k = -1)
print('Array1:\n',array1)
print('Array1 as main diagonal elements:\n', mainDiagonal)
print('Array1 as diagonal elements above main diagonal:\n', upperDiagonal)
print('Array1 as diagonal elements below main diagonal:\n', lowerDiagonal)
Output
Array1: [1 2 3] Array1 as main diagonal elements: [[1 0 0] [0 2 0] [0 0 3]] Array1 as diagonal elements above main diagonal: [[0 1 0 0] [0 0 2 0] [0 0 0 3] [0 0 0 0]] Array1 as diagonal elements below main diagonal: [[0 0 0 0] [1 0 0 0] [0 2 0 0] [0 0 3 0]]
Example 2: Extract Diagonals from a 2D Array
When a 2D array is passed to diag()
, it creates a 1D array with diagonal elements of the given array as elements.
Let us look at an example.
import numpy as np
# create a 2D array
array1 = np.arange(1, 10).reshape(3,3)
# create a 1D array with main diagonal as elements
mainDiagonal = np.diag(array1)
# create a 1D array with diagonal elements of arr1 one step above the main diagonal
upperDiagonal = np.diag(array1, k = 1)
# create a 1D array with diagonal elements of arr1 one step below the main diagonal
lowerDiagonal = np.diag(array1, k = -1)
print('Array1:\n',array1)
print('Array1\'s main diagonal elements:\n', mainDiagonal)
print('Array1\'s diagonal elements above main diagonal:\n', upperDiagona )
print('Array1\'s diagonal elements below main diagonal:\n', lowerDiagonal)
Output
Array1: [[1 2 3] [4 5 6] [7 8 9]] Array1's main diagonal elements: [1 5 9] Array1's diagonal elements above main diagonal: [2 6] Array1's diagonal elements below main diagonal: [4 8]
Related method
diagflat()
- creates a two-dimensional array with the flattened input as its diagonal.
import numpy as np
# create a 2D array
array1 = np.arange(1,5).reshape(2, 2)
# create diagonal array using diagflat()
mainDiagonal1 = np.diagflat(array1)
# create diagonal array using diag()
mainDiagonal2 = np.diag(array1.flatten())
print('Array1:\n',array1)
print('Array1\'s main diagonal elements:\n',mainDiagonal1)
print('Equivalent diag method:\n',mainDiagonal2)
Output
Array1: [[1 2] [3 4]] Array1's main diagonal elements: [[1 0 0 0] [0 2 0 0] [0 0 3 0] [0 0 0 4]] Equivalent diag method: [[1 0 0 0] [0 2 0 0] [0 0 3 0] [0 0 0 4]]
As you can see, the diagflag()
automatically flattens the 2D array and creates an array with elements of the flattened array as its diagonal.
In the case of diag()
, we manually used the flatten()
method.