Nikolaij
Nikolaij

Reputation: 321

Setting alternating off-diagonal elements in a matrix with numpy

I want to set alternating off-diagonal elements in an arbitrary NxN matrix as depicted in the figure.enter image description here

Currently I use a brute force way, which doesn’t adapt to a change of matrix size:

import numpy as np
dim = 8
size=dim**2
v=0.1
w=1

h = np.zeros(size)

valuesv = [1,8,19,26,37,44,55,62]
valuesw = [10,17,28,35,46,53]

h[valuesv] = v
h[valuesw] = w

ham=h.reshape((int(np.sqrt(size)),int(np.sqrt(size))))

Is there a more elegant way to fill the matrix? E.g. using fill_diagonal of numpy?

Upvotes: 1

Views: 896

Answers (1)

NNN
NNN

Reputation: 591

How about the following? My answer is based on this

import numpy as np

# array size
nn = 8

aa = np.zeros((nn,nn),dtype='float64')

# get indices
rows, cols = np.indices((nn,nn))

# get indices of the fist upper minor diagonal
#      k = -1 will get you the first lower minor diagonal
#      k = 2  will get you the second upper minor diagonal
row_upper = np.diag(rows, k=1)
col_upper = np.diag(cols, k=1)

# create data for the upper minor diagonal
upper_data = np.zeros(len(row_upper))
# create alternating data
upper_data[0::2] = 0.1
upper_data[1::2] = 1.0

# set data
aa[row_upper,col_upper] = upper_data

Edit

Solution which uses less memory. Eliminates creation of index matrices using np.indices.

import numpy as np
# array size
nn = 10
aa = np.zeros((nn,nn),dtype='float64')
aa.ravel()[1::(2*nn+2)]=0.1          # set upper diagonal 0.1s
aa.ravel()[2+nn::(2*nn+2)]=1.0       # set upper diagonal 1.0s
aa.ravel()[nn::(2*nn+2)]=0.1         # set lower diagonal 0.1s
aa.ravel()[2*nn+1::(2*nn+2)]=1.0     # set lower diagonal 1.0s

Upvotes: 1

Related Questions