1

I would like to set the diagonal of some numpy matrix to an arbitrary 1D array.

For example, if:

mat = np.array([[1,2],[3,4]])
diag = np.array([5,6])

then:

>>> set_diagonal(mat, diag)
>>> mat
... array([[5,2],
           [3,6]])

I could create a diagonal matrix from the diag variable, diag_fill the mat matrix to 0 and add the results. However, this seems to use a bunch of unnecessary resources if I'm trying to alter the diagonal in place.

How do I set the diagonal of a matrix in numpy?

Note: this is NOT a duplicate of this question, as they want to set the diagonal to a constant.

Him
  • 5,257
  • 3
  • 26
  • 83

3 Answers3

2

You are looking for numpy.fill_diagonal

According to the docs:

numpy.fill_diagonal(a, val, wrap=False) .
Fill the main diagonal of the given array of any dimensionality.

import numpy as np
mat = np.array([[1,2],[3,4]])

np.fill_diagonal(mat, [5,6])
print(mat)
#[[5 2]
# [3 6]]

Or

import numpy as np
mat = np.array([[1,2],[3,4]])
diag = np.array([5,6])
np.fill_diagonal(mat, diag)
print(mat)
#[[5 2]
# [3 6]]
Devesh Kumar Singh
  • 20,259
  • 5
  • 21
  • 40
  • huh. "val : scalar"... This seems to suggest that this would not work... and yet it does. So the `val` argument need not be a scalar. – Him May 01 '19 at 19:10
  • Interesting, I didn't read that fine print! But somehow it works lol! If this helped you, could you please mark my answer as accepted @Scott :) – Devesh Kumar Singh May 01 '19 at 19:14
  • You can see how it works in the [implementation](https://github.com/numpy/numpy/blob/v1.16.1/numpy/lib/index_tricks.py#L749-L861) - it ends up just assigning `val` to a slice of the array, so that works if `val` is a 1d array of the right size. Would be nice if it were documented behaviour though! – Blorgbeard May 01 '19 at 19:44
2

You can pass that 1d array to fill_diagonal.

np.fill_diagonal(mat, [5,6])
Blorgbeard
  • 101,031
  • 48
  • 228
  • 272
2

The answers using fill_diagonal are much more elegant and computationally efficient, but you can also access the diagonal indices and then change them:

mat = np.array([[1,2],[3,4]])
diag = np.array([5,6])

row, col = np.diag_indices(mat.shape[0])

mat[row, col] = diag

See here for more info:

dbouz
  • 779
  • 9
  • 14