NumPy diag()

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 be array_like)
• k (optional) - an integer number representing the diagonal to retrieve

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.