NumPy diag()

The `diag()` method either creates a new `ndarray` with the given 1D array as its diagonal elements or it extracts the diagonal from the given `ndarray`.

Example

``````import numpy as np

# create a 1D array
array1 = np.arange(3)

# create a 2D array
array2 = np.arange(9).reshape(3,3)

# create a 2D array with elements of arr1 as diagonal
diagonalArray1 = np.diag(array1)

print('Array1:\n', array1)
print('Array1 as diagonal elements:\n', diagonalArray1)

# extract diagonal elements from arr2
diagonalArray2 = np.diag(array2)

print('\nArray2\n', array2)
print('Extract diagonal of Array2\n', diagonalArray2)

'''
Array1:
[0 1 2]
Array1 as diagonal elements:
[[0 0 0]
[0 1 0]
[0 0 2]]

Array2
[[0 1 2]
[3 4 5]
[6 7 8]]
Extract diagonal of Array2
[0 4 8]
'''``````

diag() Syntax

The syntax of `diag()` is:

``numpy.diag(array, k = 0)``

diag() Arguments

The `diag()` method takes the following arguments:

• `array` - input array (can be `array_like`)
• `k` (optional) - an integer number representing the diagonal to retrieve

Note:

• By default `k` = 0 and represents the main diagonal.
• `k` > 0 represents diagonals above the main diagonal
• `k` < 0 represents diagonals below the main diagonal.

diag() Return Value

The `diag()` method either returns a new ndarray with values on the 1D array as its diagonal, or returns a 1D array containing the diagonal elements of a given ndarray.

Example 1: Create a Diagonal Array With 1D Array

When a 1D array is passed to `diag()`, it creates a diagonal array with the given array as diagonal elements.

As discussed earlier, we can use the `k` argument to control the placement of the diagonal elements in the resulting array.

Let us look at an example.

``````import numpy as np

# create a 1D array
array1 = np.arange(1, 4)

# create a 2D array with elements of array1 as the main diagonal
mainDiagonal = np.diag(array1)

# create a 2D array with elements of array1 as diagonal above the main diagonal
upperDiagonal = np.diag(array1, k = 1)

# create a 2D array with elements of array1 as  diagonal below the main diagonal
lowerDiagonal = np.diag(array1, k = -1)

print('Array1:\n',array1)
print('Array1 as main diagonal elements:\n', mainDiagonal)
print('Array1 as diagonal elements above main diagonal:\n', upperDiagonal)
print('Array1 as diagonal elements below main diagonal:\n', lowerDiagonal)``````

Output

```Array1:
[1 2 3]
Array1 as main diagonal elements:
[[1 0 0]
[0 2 0]
[0 0 3]]
Array1 as diagonal elements above main diagonal:
[[0 1 0 0]
[0 0 2 0]
[0 0 0 3]
[0 0 0 0]]
Array1 as diagonal elements below main diagonal:
[[0 0 0 0]
[1 0 0 0]
[0 2 0 0]
[0 0 3 0]]```

Example 2: Extract Diagonals from a 2D Array

When a 2D array is passed to `diag()`, it creates a 1D array with diagonal elements of the given array as elements.

Let us look at an example.

``````import numpy as np

# create a 2D array
array1 = np.arange(1, 10).reshape(3,3)

# create a 1D array with main diagonal as elements
mainDiagonal = np.diag(array1)

# create a 1D array with diagonal elements of arr1 one step above the main diagonal
upperDiagonal = np.diag(array1, k = 1)

# create a 1D array with diagonal elements of arr1 one step below the main diagonal
lowerDiagonal = np.diag(array1, k = -1)

print('Array1:\n',array1)
print('Array1\'s main diagonal elements:\n', mainDiagonal)
print('Array1\'s diagonal elements above main diagonal:\n', upperDiagona )
print('Array1\'s diagonal elements below main diagonal:\n', lowerDiagonal)``````

Output

```Array1:
[[1 2 3]
[4 5 6]
[7 8 9]]
Array1's main diagonal elements:
[1 5 9]
Array1's diagonal elements above main diagonal:
[2 6]
Array1's diagonal elements below main diagonal:
[4 8]```

Related method

`diagflat()`- creates a two-dimensional array with the flattened input as its diagonal.

``````import numpy as np

# create a 2D array
array1 = np.arange(1,5).reshape(2, 2)

# create diagonal array using diagflat()
mainDiagonal1 = np.diagflat(array1)

# create diagonal array using diag()
mainDiagonal2 = np.diag(array1.flatten())

print('Array1:\n',array1)
print('Array1\'s main diagonal elements:\n',mainDiagonal1)
print('Equivalent diag method:\n',mainDiagonal2)``````

Output

```Array1:
[[1 2]
[3 4]]
Array1's main diagonal elements:
[[1 0 0 0]
[0 2 0 0]
[0 0 3 0]
[0 0 0 4]]
Equivalent diag method:
[[1 0 0 0]
[0 2 0 0]
[0 0 3 0]
[0 0 0 4]]```

As you can see, the `diagflag()` automatically flattens the 2D array and creates an array with elements of the flattened array as its diagonal.

In the case of `diag()`, we manually used the `flatten()` method.

Our premium learning platform, created with over a decade of experience and thousands of feedbacks.

Learn and improve your coding skills like never before.

• Interactive Courses
• Certificates
• AI Help
• 2000+ Challenges