pythonnumpynumpy-ndarraydiagonal

Assign shifted diagonal values with np.diag

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.

Solution

  • 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.]]