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

Answer 2

Do you want ones on just a diagonal, or for the whole upper triangle?

There's a set of np.tri* functions to index and create upper and lower triangle arrays:

In [395]: np.triu(np.ones((5,5),int),1)
Out[395]: 
array([[0, 1, 1, 1, 1],
       [0, 0, 1, 1, 1],
       [0, 0, 0, 1, 1],
       [0, 0, 0, 0, 1],
       [0, 0, 0, 0, 0]])

I could even create a diagonal with two tri calls

In [399]: np.tril(np.triu(np.ones((5,5),int),1),1)
Out[399]: 
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]])

Not that I'd recommend this over the np.diag function. :)

or

In [404]: np.diagflat(np.ones(4,int),1)
Out[404]: 
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]])

Answer 3

You are resetting your matrix to zero on each iteration. Move the zeros call outside the loop, and it works fine:

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

Note that you also need to loop to n-1, not n-2 (moving up a row only reduces the number of 1s by 1!). Here is the output of the fixed function above:

funtest(5)

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

This matches your method, but also see the answer from Swier using numpy builtin np.diag(np.ones(n-1),1) .