Explaining numpy.diag() function (4 examples)

Introduction

Numpy is a fundamental library for scientific computing in Python. It provides a high-performance multidimensional array object and tools for working with these arrays. An essential tool in this library is the numpy.diag() function. In this tutorial, we’ll dive deep into the numpy.diag() function, exploring its capabilities with various examples ranging from basic to advanced. Understanding this function helps manipulate arrays efficiently, a critical skill in data science, machine learning, and scientific computing.

The First Glance

The numpy.diag() function serves dual purposes: it can extract a diagonal from an existing array, or it can create a diagonal array. The functionality alters based on the input and the associated parameters.

Syntax:

numpy.diag(v, k=0)

Parameters for extraction:

• v: array_like. If v is a 2-D array, return a copy of its k-th diagonal. If v is a 1-D array, return a 2-D array with v on the k-th diagonal.
• k: int, optional. Diagonal in question. The default is 0. Use k > 0 for diagonals above the main diagonal, and k < 0 for diagonals below the main diagonal.

Returns:

• output: ndarray. If v is a 2-D array, returns the k-th diagonal. If v is a 1-D array, returns a 2-D array with v on the k-th diagonal.

Extracting a Diagonal

When provided with a 2D array, numpy.diag() extracts the diagonal elements based on offset (if provided). The default behavior is to extract the main diagonal (offset=0).

import numpy as np

a = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])
diag_a = np.diag(a)
print("Diagonal elements:", diag_a)

The output will display the main diagonal elements of the array ‘a’:

Diagonal elements: [1, 5, 9]

Creating a Diagonal Array

Conversely, when a 1D array or list is provided, numpy.diag() constructs a 2D array with the provided elements on its diagonal, and zeroes elsewhere.

import numpy as np

b = np.array([1, 2, 3])
diag_b = np.diag(b)
print("Diagonal array:", diag_b)

The output will showcase a 2×2 array with the elements of ‘b’ on its diagonal:

Diagonal array:
[[1 0 0]
[0 2 0]
[0 0 3]]

Example 1: Basic Diagonal Extraction

Let’s start with a simple example of extracting the main diagonal from a square matrix.

import numpy as np

matrix = np.array([[10, 20, 30],
                   [40, 50, 60],
                   [70, 80, 90]])
main_diag = np.diag(matrix)
print("Main diagonal:", main_diag)

The output shows the main diagonal of the matrix:

Main diagonal: [10, 50, 90]

Example 2: Extracting Diagonals with Offset

The numpy.diag() function allows for extracting diagonals other than the main one by specifying the offset parameter. A positive offset moves the diagonal to the upper right, while a negative offset moves it to the lower left.

import numpy as np

offset_matrix = np.array([[10, 20, 30],
                          [40, 50, 60],
                          [70, 80, 90]])
upper_diag = np.diag(offset_matrix, k=1)
lower_diag = np.diag(offset_matrix, k=-1)
print("Upper diagonal (k=1):", upper_diag)
print("Lower diagonal (k=-1):", lower_diag)

The output will display the elements of the upper and lower diagonals based on the specified offsets:

Upper diagonal (k=1): [20, 60]

Lower diagonal (k=-1): [40, 80]

Example 3: Creating Diagonal Arrays with Offsets

Similar to extracting diagonals with an offset, you can also create diagonal arrays with elements offset from the main diagonal.

import numpy as np

offset_elements = np.array([1, 2, 3])
diagonal_with_offset = np.diag(offset_elements, k=1)
print("Diagonal array with offset (k=1):
", diagonal_with_offset)

The output will show a diagonal array where the specified elements are placed on an offset diagonal:

Diagonal array with offset (k=1):
[[0 1 0 0]
[0 0 2 0]
[0 0 0 3]
[0 0 0 0]]

Example 4: Advanced Use – Manipulating Diagonal Elements in 3D Arrays

While numpy.diag() primarily operates on 2D arrays, its functionality can be extended to manipulating diagonal elements in higher-dimensional arrays using indexing and reshaping techniques. This example showcases how to replace diagonal elements in a 3D array.

import numpy as np

three_d_array = np.zeros((3, 3, 3))
for i in range(3):
    np.fill_diagonal(three_d_array[i], i+1)
print("3D array with manipulated diagonals:\n", three_d_array)

This technique involves iterating through slices of the 3D array and using numpy.fill_diagonal() to manipulate the diagonal elements of each slice. The output demonstrates this advanced manipulation:

3D array with manipulated diagonals:
[[[1. 0. 0.]
[0. 1. 0.]
[0. 0. 1.]]
[[2. 0. 0.]
[0. 2. 0.]
[0. 0. 2.]]
[[3. 0. 0.]
[0. 3. 0.]
[0. 0. 3.]]]

Conclusion

The numpy.diag() function is a versatile tool in numpy’s arsenal, enabling both the extraction and creation of diagonal elements in arrays. From simple extractions to manipulating elements in complex structures, understanding its use cases allows for more efficient data manipulation and cleaner code. Hopefully, this tutorial has illuminated the function’s versatility and provided you with the skills to utilize numpy.diag() in your projects.