Python

# 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

1. v – the input array or array_like object.
2. 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.