Ivan Gonzalez
Ivan Gonzalez

Reputation: 576

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.

Upvotes: 3

Views: 1539

Answers (2)

Ma.Te.Pa
Ma.Te.Pa

Reputation: 290

If I have understood correctly you are slightly misinterpreting the diag(v,k) as its parameters are (sub, super)-diagonal matrix from a vector v. To achieve your desired matrix you can do the following:

 u = np.zeros(25).reshape(5,-1)
 v = np.ones(5) # a vector v, a parameter in the diag(v,k=1)
 u = u + np.diag(v[:-1], 1) # k indicates the offset (negative being sub, positive being super)
 print(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.]])

Now you can see the vector v slicing directly depending on k because the offset diagonals have fewer elements. This method can give any elements you want to an offset diagonal by creating a vector with the desired values and adding it to your zero matrix. Using this we can change the values of any offset diagonal. I hope this answers your questions and shows that we can do it for any numbers by creating a vector using full((n,m),q) or for that matter having any vector.

Upvotes: 1

Michael Szczesny
Michael Szczesny

Reputation: 5036

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

Upvotes: 3

Related Questions