python - 如何使用 numpy 创建对角矩阵？

Question

我不想修改现有数组，我想创建一个新数组。具体来说，我的矩阵应该是:

-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

我从:

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

然后尝试连接 98 个零，但这似乎不是最好的方法。任何帮助将不胜感激。

Answer 1

使用diag 制作对角矩阵:

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 具有一次设置多个对角线的方法:

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

我可以使用 np.ones(4, dtype=int) 来保留数组整数 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]])

或者使用 np.diag 使用的 flat 迭代器:

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