How can I create a diagonal matrix with numpy?
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 -2I'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.
41 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]])