CODEWITHSUNDEEP
pythonnumpy
egwene sedai
egwene sedai

Reputation: 433

numpy: fill offset diagonal with different values

I need to make a n*n matrix m whose elements follow m(i,i+1)=sqrt(i) and 0 otherwise. For example, for n=5, we should have

[0 a 0 0 0]
[0 0 b 0 0]
[0 0 0 c 0]
[0 0 0 0 d]
[0 0 0 0 0]

where {a,b,c,d}=sqrt({1,2,3,4}). Here is a solution for a constant tri-diagonal matrix, but my case is a bit more complicated than that. I know I can do that with a loop or with list comprehension, but are there other ways? n can be potentially big.

e.g. (list comprehension code)

ele=lambda i,j:sqrt(i+1) if j-i==1 else 0

[[ele(i,j) for j in range(0,6)] for i in range(0,6)]

Upvotes: 12

Views: 19410

Answers (3)

hbailo
hbailo

Reputation: 99

A bit late obviously but i want to introduce numpy diagflat method in this question. It can be done in this way:

import numpy as np

n = 5

aux = np.arange(1 , n)
aux = np.sqrt(aux)

A = np.diagflat(aux , 1)

Upvotes: 2

ybeltukov
ybeltukov

Reputation: 567

You can straightforwardly use np.diag:

>>> d = np.sqrt(1 + np.arange(4))
>>> np.diag(d, 1)
array([[ 0.        ,  1.        ,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  1.41421356,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  1.73205081,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  2.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ]])

The second argument of np.diag specifies the diagonal in question.

Upvotes: 14

Alex Riley
Alex Riley

Reputation: 176860

One way could be to create the array of zeros and then use indexing to select and fill the desired indices with the square-root values.

For example:

>>> z = np.zeros((5,5))
>>> rng = np.arange(4)
>>> z[rng, rng+1] = np.sqrt(rng+1)
>>> z
array([[ 0.        ,  1.        ,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  1.41421356,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  1.73205081,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  2.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ]])

Upvotes: 8

Related Questions