How to transform an identity matrix with for loops in python

Question

So, lets say I have an identity matrix AB and I would like to change the diagonal ones into 5 and the immediate off diagonal into 8s

    AB=np.identity(5)
    AB
    array([[1., 0., 0., 0., 0.],
           [0., 1., 0., 0., 0.],
           [0., 0., 1., 0., 0.],
           [0., 0., 0., 1., 0.],
           [0., 0., 0., 0., 1.]])
row=AB.shape[0]
col=AB.shape[1]
new = AB
for i in range (0,row):
    for j in range (0,col): 
        if AB[i,j] !=0:
            new[i,j] = 5

       ## if j+1 == 0:
           ## new[i,j+1] = 8 
            
print(new)
[[5. 0. 0. 0. 0.]
 [0. 5. 0. 0. 0.]
 [0. 0. 5. 0. 0.]
 [0. 0. 0. 5. 0.]
 [0. 0. 0. 0. 5.]]

With the code above I am able to change the 1's into 5's but I haven't been able to figure out how to change the first off diagonal to 8s on both sides of the main diagonal

you want to use `for` loop?....I mean you can transform it by writing few words? — Anurag Dabas, Feb 27 '21 at 04:28

Alain T. · Answer 1 · 2021-02-27T05:18:47.613

0

You can use the condition:

if abs(i-j)==1: new[i,j] = 8

Or, you could build the array with the values directly:

import numpy as np

N = 5
AB = np.array(([5,8,*[0]*(N-2),8]*N)[:N*N]).reshape((N,N))
print(AB)

[[5 8 0 0 0]
 [8 5 8 0 0]
 [0 8 5 8 0]
 [0 0 8 5 8]
 [0 0 0 8 5]]

Another option is to use range assignments:

import numpy as np

N = 5
AB=np.identity(N)*5
AB[range(N-1),range(1,N)]=8
AB[range(1,N),range(N-1)]=8

print(AB)
[[5. 8. 0. 0. 0.]
 [8. 5. 8. 0. 0.]
 [0. 8. 5. 8. 0.]
 [0. 0. 8. 5. 8.]
 [0. 0. 0. 8. 5.]]

edited Feb 27 '21 at 05:18

answered Feb 27 '21 at 04:29

Alain T.

40,517
4
31
51

thank you @Alain I need 8 on the first off diagonal and 0's on the rest – P20L Feb 27 '21 at 04:31

score 0 · Answer 2 · answered Feb 27 '21 at 04:31

Assuming you'll part with the idea of using a for loop, you can just write:

import numpy as np 
res = np.diag([5]*4, 0) + np.diag([8]*3, -1) + np.diag([8]*3, 1)

Here, we are telling numpy to put an array of 5's on the main diagonal and 8's on the first two off diagonals (the -1 in np.diag([8]*3, -1) says the lower left diagonal, and the +1 in np.diag([8]*3, 1) says the upper right diagonal).

[[5 8 0 0]
 [8 5 8 0]
 [0 8 5 8]
 [0 0 8 5]]

If you want to use for loops, you can do it in a clever way by paying attention to indices, however, there is a reason packages for this exist.

score 0 · Answer 3 · answered Feb 27 '21 at 04:34

To access the entries in an array you can specify the coordinates to access. In fact, you can specify all the coordinates at once and do not need to do any loops. In general, if you find yourself doing loops with numpy, you are probably doing something wrong.

import numpy as np

AB                          = np.identity(5)
AB[[0,1,2,3,4],[0,1,2,3,4]] = 5
AB[[0,1,2,3],[1,2,3,4]]     = 8
AB[[1,2,3,4],[0,1,2,3]]     = 8

print(AB)

How to transform an identity matrix with for loops in python

3 Answers3