📚 课程目录第4章案例4：神经算子学习一维扩散方程解算子

4. 案例4：神经算子学习一维扩散方程解算子

问题设定与数据构造

本节要学习的目标算子是：给定一维扩散方程的初始条件 u(x,0)=u₀(x)，输出固定时刻 T 的解 u(x,T)。方程取周期边界：

∂u/∂t = ν · ∂²u/∂x²,  x ∈ [0, 2π),  t ∈ [0, T]
u(x,0) = u₀(x),  u(x+2π,t) = u(x,t)

在周期边界下，扩散方程有解析解：对初始条件做傅里叶展开 u₀(x)=Σ aₖ e^{ikx}，则 u(x,T)=Σ aₖ e^{-νk²T} e^{ikx}。这个解析解让我们可以精确生成训练与测试数据，而不必依赖数值积分误差。

为了让算子学习有实际意义，初始条件不能是单一正弦，而应覆盖一定函数类。这里取随机有限项傅里叶级数作为训练分布：

import numpy as np

N = 128                      # 空间网格点数
x = np.linspace(0, 2*np.pi, N, endpoint=False)
nu, T = 0.1, 0.5             # 扩散系数与目标时刻
K_train = 5                  # 训练分布：|k| ≤ 5 的随机模态

def random_ic(rng, K, n_samples):
    """生成 n_samples 个随机初始条件，返回 (n_samples, N)"""
    u0 = np.zeros((n_samples, N))
    for k in range(1, K+1):
        ak = rng.normal(0, 1.0/k, size=(n_samples, 1))   # 幅值随 k 衰减
        bk = rng.normal(0, 1.0/k, size=(n_samples, 1))
        u0 += ak*np.cos(k*x) + bk*np.sin(k*x)
    return u0

def diffusion_solve(u0, nu, T):
    """周期扩散解析解：逐模态乘以衰减因子"""
    u0_hat = np.fft.rfft(u0, axis=1)
    k = np.fft.rfftfreq(N, d=1.0/N)          # 整数波数
    decay = np.exp(-nu * (k**2) * T)
    return np.fft.irfft(u0_hat * decay, n=N, axis=1)

训练集用 K_train=5 生成 2000 个样本；测试集分两组：分布内测试（同样 K=5，另取随机种子 500 个）与分布外测试（K=12，含更高频模态，500 个）。分布外测试正是检验算子泛化的关键——高频模态在 T=0.5 时衰减因子 e^{-0.1·144·0.5}≈e^{-7.2} 已很小，但初始条件的高频结构仍会通过分支网输入影响预测。

DeepONet 结构与两种容量配置

DeepONet 把解写成 u(x,T) ≈ Σⱼ bⱼ(u₀) · tⱼ(x)，其中分支网 b(u₀) 编码初始条件（在 N 个网格点上采样），主干网 t(x) 编码查询坐标。由于目标时刻固定为 T，主干网只需输入 x。两种容量配置对比如下：

配置	分支网	主干网	参数量（约）	设计意图
A 均衡	128→128→128，输出 p=64	1→128→128，输出 p=64	~5.0 万	基线，两支容量相当
B 分支偏重	128→256→256，输出 p=64	1→64→64，输出 p=64	~8.5 万	检验“输入编码不足”是否为瓶颈

注意主干网输入维度只有 1，容量需求天然较低；若分支网容量不足，高频初始条件的编码会被压缩丢失，分布外误差会显著上升。这正是本节要比较的核心。

训练实现与容量对比

下面给出可直接运行的 PyTorch 实现。关键点：分支网输入是 128 维向量，主干网输入是 1 维标量；两者输出维度必须一致（p=64），内积后得到标量预测。损失用相对 L2 误差，比 MSE 更能反映算子整体精度。

import torch, torch.nn as nn

class Branch(nn.Module):
    def __init__(self, width, p=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(N, width), nn.Tanh(),
            nn.Linear(width, width), nn.Tanh(),
            nn.Linear(width, p))
    def forward(self, u0): return self.net(u0)

class Trunk(nn.Module):
    def __init__(self, width, p=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, width), nn.Tanh(),
            nn.Linear(width, width), nn.Tanh(),
            nn.Linear(width, p))
    def forward(self, x): return self.net(x)

class DeepONet(nn.Module):
    def __init__(self, bw, tw, p=64):
        super().__init__()
        self.branch, self.trunk = Branch(bw, p), Trunk(tw, p)
    def forward(self, u0, x):
        b = self.branch(u0)                 # (B, p)
        t = self.trunk(x)                   # (N, p)
        return b @ t.T                      # (B, N)

def train(model, u0_tr, uT_tr, epochs=3000, lr=1e-3):
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    x_t = torch.tensor(x, dtype=torch.float32).unsqueeze(1)
    for ep in range(epochs):
        opt.zero_grad()
        pred = model(u0_tr, x_t)
        loss = torch.norm(pred - uT_tr) / torch.norm(uT_tr)
        loss.backward(); opt.step()
    return loss.item()

# 配置 A：分支 128 / 主干 128
# 配置 B：分支 256 / 主干 64
torch.manual_seed(0)
mA = DeepONet(128, 128); train(mA, u0_tr, uT_tr)
mB = DeepONet(256, 64);  train(mB, u0_tr, uT_tr)

训练完成后，在两组测试集上计算相对 L2 误差。预期现象（示例性判断，非实测）：配置 A 在分布内误差约 1e-3 量级，分布外（K=12）误差明显升高；配置 B 因分支网容量更大，分布外误差应低于 A，说明高频初始条件的编码能力是泛化瓶颈。若 B 的分布内误差反而更大，则可能是主干网容量不足导致坐标编码粗糙，需要回到主干网宽度上调整。

判断容量瓶颈的实用方法：固定主干网不变，只增大分支网宽度，若分布外误差下降而分布内几乎不变，说明瓶颈在输入编码；反之若两者同步下降，则瓶颈在坐标表示。每次只改一个变量，才能把误差来源归因清楚。

结果判读与验证要点

分布内 vs 分布外：分别报告两组相对 L2 误差，不要只报平均值，否则会掩盖高频泛化问题。
逐模态误差：对预测解做 FFT，比较 k=6~12 的模态幅值误差，能直接看出分支网是否丢失高频信息。
容量归因：配置 A 与 B 的参数量不同，比较时同时记录训练损失收敛曲线，避免把“训练不充分”误判为“容量不足”。
可复现性：固定随机种子生成数据与初始化，报告误差时注明是单次运行还是多次平均。

完成以上步骤后，你就得到了一个能对训练分布外初始条件给出定量泛化误差的扩散方程解算子，并且能说清分支网与主干网容量各自对误差的贡献。