Numpy: fill diagonal offset with different values
I need to create a matrix n*n
m
whose elements are followed by m(i,i+1)=sqrt(i)
and 0 otherwise. For example, for n=5
we must 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 tridiagonal matrix, but my case is a little more complicated. I know I can do this with a loop or a combo box, but are there other ways? n
could be potentially large.
eg. (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)]
source to share
One way would be to create an array of zeros and then use indexing to select and populate the desired indices with square 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. ]])
source to share
You can directly 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 np.diag
indicates the diagonal in question.
source to share