bundles / numpy 2.4.4 / numpy / diag_indices
function
numpy:diag_indices
Signature
def diag_indices(n, ndim = 2)Summary
Return the indices to access the main diagonal of an array.
Extended Summary
This returns a tuple of indices that can be used to access the main diagonal of an array a with a.ndim >= 2 dimensions and shape (n, n, ..., n). For a.ndim = 2 this is the usual diagonal, for a.ndim > 2 this is the set of indices to access a[i, i, ..., i] for i = [0..n-1].
Parameters
n: intThe size, along each dimension, of the arrays for which the returned indices can be used.
ndim: int, optionalThe number of dimensions.
Examples
import numpy as np
✓di = np.diag_indices(4) di a = np.arange(16).reshape(4, 4) a a[di] = 100 a✓
array([[100, 1, 2, 3],
[ 4, 100, 6, 7],
[ 8, 9, 100, 11],
[ 12, 13, 14, 100]])
Now, we create indices to manipulate a 3-D array:
d3 = np.diag_indices(2, 3) d3✓
(array([0, 1]), array([0, 1]), array([0, 1]))And use it to set the diagonal of an array of zeros to 1:
a = np.zeros((2, 2, 2), dtype=int) a[d3] = 1✓
a
✗array([[[1, 0],
[0, 0]],
[[0, 0],
[0, 1]]])
See also
Aliases
numpy.diag_indices