NumPy is a powerful library in Python that provides support for large, multi-dimensional arrays and matrices. One of the useful functions offered by NumPy is the diag function. The diag function can be used to extract the diagonal elements from a given matrix or to construct a diagonal matrix from an array of elements.

Extracting Diagonal Elements

To extract the diagonal elements from a matrix using the diag function, we simply pass the matrix as an argument. The diag function returns a 1-D array containing the diagonal elements.

import numpy as np

matrix = np.array([[1, 2, 3],
                   [4, 5, 6],
                   [7, 8, 9]])

diagonal = np.diag(matrix)
print(diagonal)

Output:

[1 5 9]

In the above example, we create a 3x3 matrix and use the diag function to extract the diagonal elements [1, 5, 9], which are then printed to the console.

Constructing Diagonal Matrix

The diag function can also be used to construct a diagonal matrix from an array of elements. To do this, we pass the array as an argument and specify the k parameter to specify the position of the diagonal.

import numpy as np

elements = np.array([1, 2, 3])
diagonal_matrix = np.diag(elements)
print(diagonal_matrix)

Output:

[[1 0 0]
 [0 2 0]
 [0 0 3]]

In the above example, we create a 1-D array [1, 2, 3] and construct a 3x3 diagonal matrix using the diag function. The resulting matrix is then printed, showing the diagonal elements.

The k parameter can be used to specify the position of the diagonal in the resulting matrix. A positive value of k shifts the diagonal above the main diagonal, while a negative value shifts it below. Let’s see an example:

import numpy as np

elements = np.array([1, 2, 3])
diagonal_matrix = np.diag(elements, k=1)
print(diagonal_matrix)

Output:

[[0 1 0 0]
 [0 0 2 0]
 [0 0 0 3]
 [0 0 0 0]]

In this example, we set k=1, so the diagonal elements are shifted one position above the main diagonal.

In summary, the diag function in NumPy is a versatile tool for working with diagonal elements of arrays and constructing diagonal matrices. It offers flexibility and convenience, making it a valuable function in numerical computing tasks involving matrices and arrays.