NumPy np.diag()

The NumPy diag() function lets you extract or create a new diagonal array. It is a convenient but straightforward function.

Follow along with this tutorial to explore this function further.

Function Syntax

The function syntax is as shown below:

numpy.diag(v, k=0)

The function takes just two parameters, as discussed below:

Function Parameters

v – the input array or array_like object.
k – defines the diagonal to be extracted.

NOTE: If the value of k is greater than 0, it means the diagonal above the main diagonal. If it is negative, it means the diagonal below the main diagonal.

Function Return Value

The function returns the extracted diagonal array or a newly constructed diagonal array.

Example 1

Consider the example code shown below:

import numpy as np
arr = np.arange(6).reshape(2,3)
print(arr)
print(np.diag(arr, k=0))

We use the diag() function to extract the main diagonal from the provided array in the code above.

The resulting output is shown:

[[0 1 2]
[3 4 5]]
--> extracted: [0 4]

Example 2

To extract the diagonal above the main, set the value of k as one as shown:

arr = np.arange(6).reshape(2,3)
print(f"original: {arr}")
print(f"extract: {np.diag(arr, k=1)}")

This returns:

original: [[0 1 2]
[3 4 5]]
extract: [1 5]

Example 3

If the value of k is negative, it should return:

arr = np.arange(6).reshape(2,3)
print(f"original: {arr}")
print(f"extract: {np.diag(arr, k=-1)}")

Output:

original: [[0 1 2]
[3 4 5]]
extract: [3]

Conclusion

Now you are familiar with the diag function in NumPy and how to use it to extract or construct a new diagonal array.

Thanks for reading!!