向量化LU分解求解多个b



我正在用scipy的LU分解对形状为(n,n)的二次矩阵A进行预处理,然后对形状(...,n)的多个右手边B反复求解。但scipy.linalg.lu_solve只接受b的向量,而不接受像(m,n)(k,m,n)这样的矩阵。如何包装lu_solve以处理形状为(...,n)的参数?Numpy的linalg.solve将接受多个b,但不允许单独的LU因子和求解操作。

lu_solve的文档中没有提到,但实际上b可以包含多个向量。如果A具有(n, n)的形状,则b可以具有(n, m)的形状。例如,

In [44]: A                                                                                                           
Out[44]: 
array([[ 1.01,  0.02, -0.01],
[ 0.02,  1.04, -0.02],
[-0.01, -0.02,  1.01]])
In [45]: b                                                                                                           
Out[45]: 
array([[ 0,  1,  2,  3],
[ 4,  5,  6,  7],
[ 8,  9, 10, 11]])
In [46]: lu = lu_factor(A)                                                                                           
In [47]: x = lu_solve(lu, b)                                                                                         
In [48]: x                                                                                                           
Out[48]: 
array([[ 0.        ,  0.98113208,  1.96226415,  2.94339623],
[ 4.        ,  4.96226415,  5.9245283 ,  6.88679245],
[ 8.        ,  9.01886792, 10.03773585, 11.05660377]])
In [49]: A.dot(x)                                                                                                    
Out[49]: 
array([[ 0.,  1.,  2.,  3.],
[ 4.,  5.,  6.,  7.],
[ 8.,  9., 10., 11.]])

更高维度的b必须具有形状(n, ...)。请注意,对于二维以上的形状,用A.dot(x)测试结果将不起作用,因为x的形状与NumPy的矩阵乘法不兼容。例如,这里的B具有形状(3, 2, 5):

In [40]: A                                                                      
Out[40]: 
array([[ 1.01,  0.02, -0.01],
[ 0.02,  1.04, -0.02],
[-0.01, -0.02,  1.01]])
In [41]: B = np.random.rand(3, 2, 5)                                            
In [42]: lu = lu_factor(A)                                                      
In [43]: x = lu_solve(lu, B)                                                    
In [44]: x.shape                                                                
Out[44]: (3, 2, 5)
In [45]: xx = np.moveaxis(x, 0, 1)                                              
In [46]: np.allclose(A.dot(xx), B)                                              
Out[46]: True