Python code to produce diagonal matrix with 1 above the main diagonal

Question

I would like to write a python function that produce a zero matrix of size n by n that has 1 for all the elements above the main diagonal.

Here is my code:

def funtest(n):
    for i in range(0,n-2):
        s = (n,n)
        Y = zeros(s)
        Z = Y
        Z[i,i+1]=1          
return Z

But the result only give 1 in the (n-1, n-2) element of the matrix.

I am stuck and I really think my code is correct and have no clue where is the mistake. And how can I fix it? Can anybody please help?

Thanks.

Answer 1

The numpy.diag function can do exactly this:

import numpy as np
print( np.diag(np.ones(4), 1) )

With the second argument (the 1) the offset for the diagonal. It gives:

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