NumPy is a powerful Python library, which supports many mathematical functions that can be applied to multi-dimensional arrays.

In this short tutorial, you will learn how to create diagonal arrays with NumPy.

Working with Diagonal Arrays in NumPy

NumPy has two built-in functions, np.diag and np.diagflat, which may be used to create diagonal arrays. Each function has its own distinct purpose:

• The np.diag function outputs a diagonal array or retrieves the diagonal elements of an input array, depending on whether its input is one-dimensional (1D) or two-dimensional (2D), respectively.
• The np.diagflat function outputs a diagonal array irrespective of whether its input is 1D or 2D.

Let’s have a look at these two functions in greater detail.

Using the diag Function in Numpy

The built-in np.diag function in Numpy takes takes two input arguments: v of type array_like and an optional integer, k.

Let’s have a look at the role of v first.

The input argument, v, can be a 1D or a 2D array. Depending upon the dimensionality of the input, the output that is produced by the np.diag function is as follows:

• If v is a 1D array, the np.diag function returns a 2D array with the values in v on its diagonal.
• If v is a 2D array, the np.diag function extracts its diagonal elements and returns them (for this first example we will retain the default value of k = 0, but we shall be returning to exploring the role of k later).

Let’s try setting v to a 1D array first:

# Import NumPy package
import numpy as np

# Define 1D array
v = np.array([1, 2, 3, 4])

# Create diagonal array
diag_array = np.diag(v)

# Print diagonal array
print('Diagonal array: \n', diag_array)

We find that the np.diag function outputs a 4×4 array with the values in v on its main diagonal:

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

If, alternatively, the input to the function is a 2D array:

# Define 2D array
v = np.array([[1, 0, 1], [3, 2, 3], [4, 3, 4]])

# Print 2D array
print('2D input array: \n', v)

# Extract diagonal elements
diag_elements = np.diag(v)

# Print diagonal elements
print('\n Diagonal elements: \n', diag_elements)

We find that the function extracts the diagonal elements of its 2D input and returns them:

2D input array: 
 [[1 0 1]
 [3 2 3]
 [4 3 4]]

 Diagonal elements: 
 [1 2 4]

Now let’s go back to exploring the role of optional input argument, k, by first considering a 1D input to the function.

When set to a value that is lower or higher than its default 0, the parameter k determines the direction to which the array’s diagonal is offset from the main diagonal, and by how much. For example, let’s try of value of k = 1:

# Define 1D array
v = np.array([1, 2, 3, 4])

# Create offset diagonal array
off_diag_array = np.diag(v, k=1)

# Print offset diagonal array
print('Offset diagonal array: \n', off_diag_array)

In this case, the np.diag function produces a 5×5 array with its diagonal elements shifted rightwards to (or above) the main diagonal by a step size of 1 (corresponding to a value of k = 1):

Offset diagonal array: 
 [[0 1 0 0 0]
 [0 0 2 0 0]
 [0 0 0 3 0]
 [0 0 0 0 4]
 [0 0 0 0 0]]

If we had to try setting a value of k = -1, we would find that the diagonal elements now shift leftwards to (or below) the main diagonal, also by a step size of 1:

Offset diagonal array: 
 [[0 0 0 0 0]
 [1 0 0 0 0]
 [0 2 0 0 0]
 [0 0 3 0 0]
 [0 0 0 4 0]]

If the input to the np.diag function is rather in 2D:

# Define 2D array
v = np.array([[1, 0, 1], [3, 2, 3], [4, 3, 4]])

# Print 2D array
print('2D input array: \n', v)

# Extract offset diagonal elements
off_diag_elements = np.diag(v, k=1)

# Print offset diagonal elements
print('\n Diagonal elements: \n', off_diag_elements)

Specifying a value for k that is not its default 0, results in the np.diag function extracting the diagonal elements of the input array that are offset from its main diagonal according to the value of k:

2D input array: 
 [[1 0 1]
 [3 2 3]
 [4 3 4]]

 Diagonal elements: 
 [0 3]

Using the diagflat Function in Numpy

Another option in creating diagonal arrays with NumPy is to use the np.diagflat function. Similar to the np.diag function reviewed earlier, the np.diagflat function takes an input argument v of type array_like, and an optional integer argument, k.

The value of the input argument k determines the direction and step size of the offset that is applied to the diagonal elements, away from the main diagonal of the array. This works identically to what we have seen for the np.diag function.

What the np.diagflat function does differently, on the other hand, is that it produces a 2D diagonal array regardless of whether the v input is 1D or 2D. Indeed, if the input array is 2D, the np.diagflat function flattens it out and places the elements on the diagonal of its output array.

Let’s test this out:

# Define 1D array
v_1D = np.array([1, 2, 3, 4])

# Define 2D array
v_2D = np.array([[1, 2], [3, 4]])

# Create diagonal array from 1D input
diag_array_1D = np.diagflat(v_1D)

# Create diagonal array from 2D input
diag_array_2D = np.diagflat(v_2D)

# Print diagonal arrays
print('Diagonal array for 1D input: \n', diag_array_1D)
print('\n Diagonal array for 2D input: \n', diag_array_2D)

We find that, when using the np.diagflat function, the code above produces identical results regardless of the dimensionality of the input:

Diagonal array for 1D input: 
 [[1 0 0 0]
 [0 2 0 0]
 [0 0 3 0]
 [0 0 0 4]]

 Diagonal array for 2D input: 
 [[1 0 0 0]
 [0 2 0 0]
 [0 0 3 0]
 [0 0 0 4]]