python - 优化三对角系数矩阵的 A*x = B 解

Question

我有一个形式为 A*x = B 的方程组，其中 [A] 是一个三对角系数矩阵。使用 Numpy 求解器 numpy.linalg.solve 我可以求解 x 的方程组。

请参阅下面的示例，了解我如何开发三对角线 [A] 矩阵。 {B} 向量，求解 x:

# Solve system of equations with a tridiagonal coefficient matrix
# uses numpy.linalg.solve

# use Python 3 print function
from __future__ import print_function
from __future__ import division

# modules
import numpy as np
import time

ti = time.clock()

#---- Build [A] array and {B} column vector

m = 1000   # size of array, make this 8000 to see time benefits

A = np.zeros((m, m))     # pre-allocate [A] array
B = np.zeros((m, 1))     # pre-allocate {B} column vector

A[0, 0] = 1
A[0, 1] = 2
B[0, 0] = 1

for i in range(1, m-1):
    A[i, i-1] = 7   # node-1
    A[i, i] = 8     # node
    A[i, i+1] = 9   # node+1
    B[i, 0] = 2

A[m-1, m-2] = 3
A[m-1, m-1] = 4
B[m-1, 0] = 3

print('A \n', A)
print('B \n', B)

#---- Solve using numpy.linalg.solve

x = np.linalg.solve(A, B)     # solve A*x = B for x

print('x \n', x)

#---- Elapsed time for each approach

print('NUMPY time', time.clock()-ti, 'seconds')

所以我的问题与上面例子的两个部分有关:

1. 由于我正在处理 [A] 的三对角矩阵，也称为带状矩阵，是否有更有效的方法来求解方程组而不是使用 numpy.linalg .解决?
2. 此外，是否有更好的方法来创建三对角矩阵而不是使用 for-loop？

根据 time.clock() 函数，以上示例在 Linux 上运行大约 0.08 秒。

numpy.linalg.solve 函数工作正常，但我试图找到一种利用 [A] 的三对角形式的方法，希望进一步加快求解速度，然后将该方法应用于更复杂的示例。

Answer 1

有两个直接的性能改进 (1) 不使用循环，(2) 使用 scipy.linalg.solve_banded()。

我会把代码写得更像

import scipy.linalg as la

# Create arrays and set values
ab = np.zeros((3,m))
b = 2*ones(m)
ab[0] = 9
ab[1] = 8
ab[2] = 7

# Fix end points
ab[0,1] = 2
ab[1,0] = 1
ab[1,-1] = 4
ab[2,-2] = 3
b[0] = 1
b[-1] = 3

return la.solve_banded ((1,1),ab,b)

可能有更优雅的方法来构造矩阵，但这行得通。

在 ipython 中使用 %timeit 原始代码在 m=1000 时花费了 112 毫秒。当 m=10,000 时，此代码耗时 2.94 毫秒，问题大了一个数量级，但仍然快了将近两个数量级!我没有耐心等待m = 10,000 的原始代码。原来大部分时间可能是在构造数组，这个我没测试。无论如何，对于大型数组，仅存储矩阵的非零值会更有效。