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

Question

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

请参阅下面的示例，了解我如何开发三对角线 [A] martix。 {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.solve 更有效的方法来求解方程组？
2. 另外，有没有更好的方法来创建三对角矩阵而不是使用for-loop？

根据time.clock()函数，上面的例子在大约0.08 seconds的Linux上运行。

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

Answer 1

有一个名为scipy.sparse.dia_matrix 的scipy.sparse 矩阵类型可以很好地捕获矩阵的结构（它将存储3 个数组，在“位置”0（对角线）、1（上）和-1（下）） .使用这种类型的矩阵，您可以尝试scipy.sparse.linalg.lsqr 进行求解。如果你的问题有一个精确的解决方案，它就会被找到，否则它会在最小二乘意义上找到解决方案。

from scipy import sparse
A_sparse = sparse.dia_matrix(A)
ret_values = sparse.linalg.lsqr(A_sparse, C)
x = ret_values[0]

但是，就利用三对角结构而言，这可能不是完全最优的，可能有一种理论上的方法可以加快速度。此转换为您做的是将矩阵乘法费用减少到基本：仅使用 3 个波段。这与迭代求解器 lsqr 结合应该已经产生了加速。

注意：我不是在提议scipy.sparse.linalg.spsolve，因为它会将您的矩阵转换为csr 格式。但是，将lsqr 替换为spsolve 值得一试，尤其是因为spsolve 可以绑定UMFPACK，请参阅相关doc on spsolve。另外，看看this stackoverflow question and answer relating to UMFPACK

可能会很有趣

Answer 2

你可以使用scipy.linalg.solveh_banded。

编辑：您不能使用上述方法，因为您的矩阵不是对称的，我认为是。但是，正如上面评论中提到的，Thomas 算法非常适合此

a =       [7] * ( m - 2 ) + [3]
b = [1] + [8] * ( m - 2 ) + [4]
c = [2] + [9] * ( m - 2 )
d = [1] + [2] * ( m - 2 ) + [3]

# This is taken directly from the Wikipedia page also cited above
# this overwrites b and d
def TDMASolve(a, b, c, d):
    n = len(d) # n is the numbers of rows, a and c has length n-1
    for i in xrange(n-1):
        d[i+1] -= 1. * d[i] * a[i] / b[i]
        b[i+1] -= 1. * c[i] * a[i] / b[i]
    for i in reversed(xrange(n-1)):
        d[i] -= d[i+1] * c[i] / b[i+1]
    return [d[i] / b[i] for i in xrange(n)]

这段代码没有优化，也没有使用np，但如果我（或这里的任何其他好人）有时间，我会编辑它，让它做这些事情。对于 m=10000，当前时间约为 10 ms。

Answer 3

有两个直接的性能改进（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 的原始代码。原来大部分时间可能是在构造数组，这个我没有测试。无论如何，对于大型数组，只存储矩阵的非零值会更有效。

Answer 4

这可能会有所帮助 有一个函数 create_tridiagonal 将创建三对角矩阵。还有另一个函数可以根据 SciPy solve_banded 函数的要求将矩阵转换为对角有序形式。

import numpy as np    

def lu_decomp3(a):
    """
    c,d,e = lu_decomp3(a).
    LU decomposition of tridiagonal matrix a = [c\d\e]. On output
    {c},{d} and {e} are the diagonals of the decomposed matrix a.
    """
    n = np.diagonal(a).size
    assert(np.all(a.shape ==(n,n))) # check if square matrix

    d = np.copy(np.diagonal(a)) # without copy (assignment destination is read-only) error is raised 
    e = np.copy(np.diagonal(a, 1))
    c = np.copy(np.diagonal(a, -1)) 

    for k in range(1,n):
        lam = c[k-1]/d[k-1]
        d[k] = d[k] - lam*e[k-1]
        c[k-1] = lam
    return c,d,e

def lu_solve3(c,d,e,b):
    """
    x = lu_solve(c,d,e,b).
    Solves [c\d\e]{x} = {b}, where {c}, {d} and {e} are the
    vectors returned from lu_decomp3.
    """
    n = len(d)
    y = np.zeros_like(b)

    y[0] = b[0]
    for k in range(1,n): 
        y[k] = b[k] - c[k-1]*y[k-1]

    x = np.zeros_like(b)
    x[n-1] = y[n-1]/d[n-1] # there is no x[n] out of range
    for k in range(n-2,-1,-1):
        x[k] = (y[k] - e[k]*x[k+1])/d[k]
    return x

from scipy.sparse import diags
def create_tridiagonal(size = 4):
    diag = np.random.randn(size)*100
    diag_pos1 = np.random.randn(size-1)*10
    diag_neg1 = np.random.randn(size-1)*10

    a = diags([diag_neg1, diag, diag_pos1], offsets=[-1, 0, 1],shape=(size,size)).todense()
    return a

a = create_tridiagonal(4)
b = np.random.randn(4)*10

print('matrix a is\n = {} \n\n and vector b is \n {}'.format(a, b))

c, d, e = lu_decomp3(a)
x = lu_solve3(c, d, e, b)

print("x from our function is {}".format(x))

print("check is answer correct ({})".format(np.allclose(np.dot(a, x), b)))


## Test Scipy
from scipy.linalg import solve_banded

def diagonal_form(a, upper = 1, lower= 1):
    """
    a is a numpy square matrix
    this function converts a square matrix to diagonal ordered form
    returned matrix in ab shape which can be used directly for scipy.linalg.solve_banded
    """
    n = a.shape[1]
    assert(np.all(a.shape ==(n,n)))

    ab = np.zeros((2*n-1, n))

    for i in range(n):
        ab[i,(n-1)-i:] = np.diagonal(a,(n-1)-i)

    for i in range(n-1): 
        ab[(2*n-2)-i,:i+1] = np.diagonal(a,i-(n-1))


    mid_row_inx = int(ab.shape[0]/2)
    upper_rows = [mid_row_inx - i for i in range(1, upper+1)]
    upper_rows.reverse()
    upper_rows.append(mid_row_inx)
    lower_rows = [mid_row_inx + i for i in range(1, lower+1)]
    keep_rows = upper_rows+lower_rows
    ab = ab[keep_rows,:]


    return ab

ab = diagonal_form(a, upper=1, lower=1) # for tridiagonal matrix upper and lower = 1

x_sp = solve_banded((1,1), ab, b)
print("is our answer the same as scipy answer ({})".format(np.allclose(x, x_sp)))