NumPy diag()

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) - the diagonal in question(can be integer)

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 see an example.

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 see an example.

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.

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.