如何用solve_ivp解谱方法求解偏微分方程？

import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

# 参数设置
L = 10  # 空间长度
T = 10  # 总时间
Nx = 50  # 空间离散点数
Nt = 1000  # 时间步数
alpha = 0.1  # 热扩散系数

# 空间网格
x = np.linspace(0, L, Nx, endpoint=False)
dx = x[1] - x[0]

# 时间网格
t_span = (0, T)
t_eval = np.linspace(t_span[0], t_span[1], Nt)

# 初始条件
u0 = np.sin(np.pi * x / L)

# 谱方法的系数矩阵
k = np.fft.fftfreq(Nx, d=dx) * 2 * np.pi
k2 = k**2
k2[Nx//2] = 1  # 避免除以零

# 定义 ODE 系统
def heat_eq(t, u):
    u_hat = np.fft.fft(u)
    dudt_hat = -alpha * k2 * u_hat
    return np.real(np.fft.ifft(dudt_hat))

# 使用 solve_ivp 求解
sol = solve_ivp(heat_eq, t_span, u0, t_eval=t_eval)

# 绘制结果
for i in range(0, sol.y.shape[1], 100):
    plt.plot(x, sol.y[:, i], label=f't={sol.t[i]:.2f}')

plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.legend()
plt.show()