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

import numpy as np from scipy.linalg import toeplitz

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

def pseudospectral2nd(order, fun, a, b, alpha, beta, N): # order: 方程阶数 # fun: 方程组的右端项 # a, b: 区间端点 # alpha, beta: 边界条件 # 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])