伪谱法求解二阶微分方程组的 Python 代码

import numpy as np from scipy.linalg import toeplitz

定义伪谱法求解二阶微分方程组的函数

def pseudospectral2nd(order, fun, a, b, alpha, beta, N): ''' 使用伪谱法求解二阶微分方程组

参数
----
order : int
    方程阶数
fun : function
    方程组的右端项，函数形式为 fun(x, u)，其中 x 为自变量，u 为解向量
a, b : float
    区间端点
alpha, beta : function
    边界条件，函数形式为 alpha(x) 和 beta(x)
N : int
    伪谱法的节点数

返回值
----
u : ndarray
    求解得到的函数值，形状为 (order, N)
x : ndarray
    伪谱法的节点，形状为 (N,)
'''

# 计算伪谱法的节点和权重
x, w = np.polynomial.legendre.leggauss(N)
x = (b - a)/2 * x + (b + a)/2
w = (b - a)/2 * w

# 构造初始猜测函数
u = np.zeros((order, N))
for i in range(N):
    u[0, i] = alpha(x[i])
    u[1, i] = (beta(x[i]) - alpha(x[i])) / (b - a) * 2 * (x[i] - a) + alpha(x[i])

# 构造差分矩阵
D = np.polynomial.legendre.legder(np.polynomial.legendre.legder(np.eye(N)))
D = D / ((b - a)/2)**2

# 构造对角矩阵
diag = np.eye(N)

# 构造线性方程组的系数矩阵
A = np.zeros((order*N, order*N))
A[0:N, :] = np.hstack((diag, np.zeros((N, (order-1)*N))))
for i in range(1, order):
    A[i*N:(i+1)*N, (i-1)*N:i*N] = -D
    A[i*N:(i+1)*N, i*N:(i+1)*N] = 2*diag
    A[i*N:(i+1)*N, i*N+N:(i+2)*N] = -D

# 构造线性方程组的右端项
f = np.zeros((order, N))
for i in range(N):
    f[0, i] = fun(x[i], u[:, i])[0]
    f[1, i] = fun(x[i], u[:, i])[1]
b = np.concatenate(f)

# 解线性方程组
U = np.linalg.solve(A, b)
U = U.reshape((order, N))

# 更新猜测函数
u = u + U

# 返回伪谱法求解得到的函数值和节点
return u, x

示例：使用伪谱法求解二阶常微分方程组y'' + 2y' + y = x, y(0) = 0, y(1) = 1

def fun(x, u): return np.array([x - 2u[1] - u[0], u[0] - 2u[1] + x])

def alpha(x): return 0

def beta(x): return 1

u, x = pseudospectral2nd(2, fun, 0, 1, alpha, beta, 20) print(u[0]) print(u[1])