python - Make special diagonal matrix in Numpy

Question

I am trying to make a numpy array that looks like this:

[a b c       ]
[  a b c     ]
[    a b c   ]
[      a b c ] 

So this involves updating the main diagonal and the two diagonals above it.

What would be an efficient way of doing this?

Answer 1

You can use np.indices to get the indices of your array and then assign the values where you want.

a = np.zeros((5,10))
i,j = np.indices(a.shape)

i,j are the line and column indices, respectively.

a[i==j] = 1.
a[i==j-1] = 2.
a[i==j-2] = 3.

will result in:

array([[ 1.,  2.,  3.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  1.,  2.,  3.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  1.,  2.,  3.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  1.,  2.,  3.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  1.,  2.,  3.,  0.,  0.,  0.]])