CODEWITHSUNDEEP
pythonnumpy
Shamoon
Shamoon

Reputation: 43501

How can I create a diagonal matrix with numpy?

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.

Upvotes: 8

Views: 39084

Answers (1)

hpaulj
hpaulj

Reputation: 231385

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]])

Upvotes: 12

Related Questions