2

I'm trying to assign values to multiple diagonals of a matrix. For example, I have this matrix:

>>> u = np.zeros(25).reshape(5, 5)
>>> u
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]])

I want to assign a value to the $k$-ith diagonal above. For example, if $k=1$, I want the diagonal above the main diagonal. I tried to accomplish this by using np.diag like this np.diag(u, k=1) = 1 which I want to result in the following:

>>> u
array([[0., 1., 0., 0., 0.],
       [0., 0., 1., 0., 0.],
       [0., 0., 0., 1., 0.],
       [0., 0., 0., 0., 1.],
       [0., 0., 0., 0., 0.]])

The thing is this throws me a SyntaxError: can't assign to function call as this goes again Python. But np.diag returns a reference to the original matrix as you can see:

>>> np.may_share_memory(np.diag(u, k=1), u)
True

How can I do this? Thank you in advance.

Ivan Gonzalez
  • 446
  • 1
  • 5
  • 14
  • Regardless of what it returns, python does not allow you assign to the call. That just basic python. – hpaulj Dec 15 '20 at 03:03

1 Answers1

3

You can use

u[np.eye(len(u), k=1, dtype='bool')] = 1
print(u)

Out:

[[0. 1. 0. 0. 0.]
 [0. 0. 1. 0. 0.]
 [0. 0. 0. 1. 0.]
 [0. 0. 0. 0. 1.]
 [0. 0. 0. 0. 0.]]
Michael Szczesny
  • 4,911
  • 5
  • 15
  • 32
  • 1
    Awesome! It should be an option to do `np.fill_diagonal(u, 1, k=1)`. Maybe I can make an issue on GitHub. What do you think? – Ivan Gonzalez Dec 15 '20 at 03:11
  • 1
    I know. That was my first thought, too. But no `k` argument for `fill_diagonal`. I don't think that would be a priority for the developers. My solution feels like a hack, perhaps someone will post a nicer approach. – Michael Szczesny Dec 15 '20 at 03:16