Velvet Star Monitor

Standout celebrity highlights with iconic style.

general

How can I create a diagonal matrix with numpy?

Writer Matthew Barrera

I do not want to modify an existing array, I want to create a new array. Specifically, my matrix should be:

-2 1 0 0 0 0 ... 0 1 -2 1 0 0 0 ... 0 0 1 -2 1 0 0 ... 0
... 0 ..........1 -2 1 0 ..........0 1 -2

I'm starting with:

 self.A = np.array([-2, 1])

then trying to concatenate 98 zeroes, but this seems like it's not the best way. Any help would be greatly appreciated.

4

1 Answer

Using diag to make a diagonal matrix:

In [140]: np.diag(np.full(5,-2))
Out[140]:
array([[-2, 0, 0, 0, 0], [ 0, -2, 0, 0, 0], [ 0, 0, -2, 0, 0], [ 0, 0, 0, -2, 0], [ 0, 0, 0, 0, -2]])
In [141]: np.diag(np.ones(4),1)
Out[141]:
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.]])
In [142]: np.diag(np.full(5,-2))+np.diag(np.ones(4),1)+np.diag(np.ones(4),-1)
Out[142]:
array([[-2., 1., 0., 0., 0.], [ 1., -2., 1., 0., 0.], [ 0., 1., -2., 1., 0.], [ 0., 0., 1., -2., 1.], [ 0., 0., 0., 1., -2.]])

scipy.sparse has methods of setting several diagonals at once:

In [143]: from scipy import sparse
In [144]: sparse.diags?
In [145]: sparse.diags([np.full(5,-2),np.ones(4),np.ones(4)],[0,-1,1])
Out[145]:
<5x5 sparse matrix of type '<class 'numpy.float64'>' with 13 stored elements (3 diagonals) in DIAgonal format>
In [147]: sparse.diags([np.full(5,-2),np.ones(4),np.ones(4)],[0,-1,1]).A
Out[147]:
array([[-2., 1., 0., 0., 0.], [ 1., -2., 1., 0., 0.], [ 0., 1., -2., 1., 0.], [ 0., 0., 1., -2., 1.], [ 0., 0., 0., 1., -2.]])

I could use np.ones(4, dtype=int) to keep the array integer dtype.

In [148]: A = np.zeros((5,5),int)
In [149]: A[range(5),range(5)]=-2
In [150]: A[range(4),range(1,5)]=1
In [151]: A[range(1,5),range(4)]=1
In [152]: A
Out[152]:
array([[-2, 1, 0, 0, 0], [ 1, -2, 1, 0, 0], [ 0, 1, -2, 1, 0], [ 0, 0, 1, -2, 1], [ 0, 0, 0, 1, -2]])

Or using the flat iterator that np.diag uses:

In [163]: A = np.zeros((5,5),int)
In [164]: A.flat[0::6] = -2
In [165]: A.flat[1::6] = 1
In [166]: A.flat[5::6] = 1
In [167]: A
Out[167]:
array([[-2, 1, 0, 0, 0], [ 1, -2, 1, 0, 0], [ 0, 1, -2, 1, 0], [ 0, 0, 1, -2, 1], [ 0, 0, 0, 1, -2]])

Your Answer

Sign up or log in

Sign up using Google Sign up using Facebook Sign up using Email and Password

Post as a guest

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy