NumPy diag()

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 be array_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 diagonal
    • k < 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.

Our premium learning platform, created with over a decade of experience and thousands of feedbacks.

Learn and improve your coding skills like never before.

Try Programiz PRO
  • Interactive Courses
  • Certificates
  • AI Help
  • 2000+ Challenges