1. 一维热传导方程的后验不确定性量化

任务目标

构建一维热传导方程 \(u_t = \alpha u_{xx}\) 的贝叶斯PINN（B-PINN）模型。使用变分推断近似网络权重的后验分布，从而获得温度场预测的均值与方差，并通过与精确解对比评估不确定性校准质量。

问题设定与精确解

求解域：\(x \in [0, 1]\)，\(t \in [0, 1]\)，热扩散系数 \(\alpha = 0.4\)。边界与初始条件由精确解 \(u_{exact}(x,t) = e^{-t}\sin(\pi x)\) 导出（该函数满足方程与齐次边界条件）。训练数据从精确解中采样并添加高斯噪声（标准差 \(\sigma_n = 0.05\)），用于模拟带噪声的观测。

依赖与环境

需要：Python 3.9+，torch，matplotlib（可选，用于绘图）。安装命令：pip install torch matplotlib。本实现使用PyTorch的AutoGrad计算PINN残差，使用torch.distributions.Normal构建变分后验。

核心实现

以下代码完整实现了B-PINN：网络权重服从高斯变分分布，通过最小化负ELBO（均方误差损失 + KL散度）训练。训练后，用后验均值作为预测均值，后验方差传播得到预测方差。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Normal
import numpy as np

torch.manual_seed(42)
np.random.seed(42)

# ---------- 精确解与数据生成 ----------
alpha = 0.4
def u_exact(x, t):
    return torch.exp(-t) * torch.sin(np.pi * x)

# 训练点：边界 + 内部 + 初始条件
N_b = 40   # 边界点（每边）
N_i = 60   # 内部配点（用于残差）
N_u = 80   # 带噪声的观测点

# 观测数据（从精确解采样加噪声）
x_u = torch.rand(N_u, 1) * 0.8 + 0.1  # 避开边界附近
t_u = torch.rand(N_u, 1) * 0.9 + 0.05
u_obs = u_exact(x_u, t_u) + 0.05 * torch.randn(N_u, 1)

# 边界与初始配点（用于硬约束或软约束，这里用软约束）
# 左边界 x=0, 右边界 x=1
x_b = torch.cat([torch.zeros(N_b, 1), torch.ones(N_b, 1)])
t_b = torch.rand(2 * N_b, 1)
u_b = u_exact(x_b, t_b)  # 边界精确值（已知为0，但这里用精确值做软约束）

# 初始条件 t=0
x_ic = torch.rand(N_i, 1)
t_ic = torch.zeros(N_i, 1)
u_ic = u_exact(x_ic, t_ic)

# 内部配点（用于PDE残差）
x_r = torch.rand(N_i, 1)
t_r = torch.rand(N_i, 1)

# ---------- 贝叶斯PINN网络 ----------
class BayesianPINN(nn.Module):
    def __init__(self, n_input=2, n_hidden=50, n_output=1):
        super().__init__()
        self.n_hidden = n_hidden
        # 变分参数：权重和偏置的均值与log标准差
        self.fc1_mean = nn.Linear(n_input, n_hidden)
        self.fc1_logstd = nn.Linear(n_input, n_hidden)
        self.fc2_mean = nn.Linear(n_hidden, n_hidden)
        self.fc2_logstd = nn.Linear(n_hidden, n_hidden)
        self.out_mean = nn.Linear(n_hidden, n_output)
        self.out_logstd = nn.Linear(n_hidden, n_output)
        self._init_weights()

    def _init_weights(self):
        for m in [self.fc1_mean, self.fc2_mean, self.out_mean]:
            nn.init.xavier_uniform_(m.weight)
            nn.init.zeros_(m.bias)
        for m in [self.fc1_logstd, self.fc2_logstd, self.out_logstd]:
            nn.init.zeros_(m.weight)
            nn.init.constant_(m.bias, -3)  # 初始小标准差

    def forward(self, x, t, n_samples=1):
        """返回预测均值与标准差（形状: [n_samples, batch, 1]）"""
        xt = torch.cat([x, t], dim=1)
        means = []
        stds = []
        for _ in range(n_samples):
            # 采样权重
            w1 = self.fc1_mean.weight + torch.exp(self.fc1_logstd.weight) * torch.randn_like(self.fc1_mean.weight)
            b1 = self.fc1_mean.bias + torch.exp(self.fc1_logstd.bias) * torch.randn_like(self.fc1_mean.bias)
            w2 = self.fc2_mean.weight + torch.exp(self.fc2_logstd.weight) * torch.randn_like(self.fc2_mean.weight)
            b2 = self.fc2_mean.bias + torch.exp(self.fc2_logstd.bias) * torch.randn_like(self.fc2_mean.bias)
            wo = self.out_mean.weight + torch.exp(self.out_logstd.weight) * torch.randn_like(self.out_mean.weight)
            bo = self.out_mean.bias + torch.exp(self.out_logstd.bias) * torch.randn_like(self.out_mean.bias)

            h = torch.tanh(xt @ w1.T + b1)
            h = torch.tanh(h @ w2.T + b2)
            pred = h @ wo.T + bo
            means.append(pred)
            # 预测标准差（从输出层方差传播，简化处理）
            # 这里用输出层权重的方差近似预测不确定性
            out_var = torch.sum(torch.exp(2 * self.out_logstd.weight), dim=1) + torch.exp(2 * self.out_logstd.bias)
            stds.append(torch.sqrt(out_var).expand_as(pred))
        return torch.stack(means), torch.stack(stds)

# ---------- 损失函数（负ELBO） ----------
def kl_divergence(model):
    """计算变分后验与标准正态先验的KL散度（近似）"""
    kl = 0
    for layer_mean, layer_logstd in [(model.fc1_mean, model.fc1_logstd),
                                     (model.fc2_mean, model.fc2_logstd),
                                     (model.out_mean, model.out_logstd)]:
        mu = layer_mean.weight
        logstd = layer_logstd.weight
        kl += torch.sum(logstd + 0.5 * (mu**2 + torch.exp(2*logstd)) - 0.5 - logstd)
        # 偏置
        mu_b = layer_mean.bias
        logstd_b = layer_logstd.bias
        kl += torch.sum(logstd_b + 0.5 * (mu_b**2 + torch.exp(2*logstd_b)) - 0.5 - logstd_b)
    return kl

def pde_residual(model, x, t):
    """计算PINN残差 u_t - alpha*u_xx"""
    x.requires_grad_(True)
    t.requires_grad_(True)
    # 使用后验均值网络（前向一次）
    xt = torch.cat([x, t], dim=1)
    h = torch.tanh(xt @ model.fc1_mean.weight.T + model.fc1_mean.bias)
    h = torch.tanh(h @ model.fc2_mean.weight.T + model.fc2_mean.bias)
    u = h @ model.out_mean.weight.T + model.out_mean.bias
    u_t = torch.autograd.grad(u, t, grad_outputs=torch.ones_like(u), create_graph=True)[0]
    u_x = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u), create_graph=True)[0]
    u_xx = torch.autograd.grad(u_x, x, grad_outputs=torch.ones_like(u_x), create_graph=True)[0]
    return u_t - alpha * u_xx

# ---------- 训练 ----------
model = BayesianPINN()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
n_epochs = 3000
kl_weight = 1e-4  # KL权重（需调参）

for epoch in range(n_epochs):
    optimizer.zero_grad()
    # 数据损失（观测点）
    pred_mean, pred_std = model(x_u, t_u, n_samples=1)
    data_loss = torch.mean((pred_mean[0] - u_obs)**2) / (2 * 0.05**2)
    # 边界与初始条件损失
    pred_b, _ = model(x_b, t_b, n_samples=1)
    loss_b = torch.mean((pred_b[0] - u_b)**2)
    pred_ic, _ = model(x_ic, t_ic, n_samples=1)
    loss_ic = torch.mean((pred_ic[0] - u_ic)**2)
    # PDE残差损失
    res = pde_residual(model, x_r, t_r)
    loss_pde = torch.mean(res**2)
    # KL散度
    kl = kl_divergence(model)
    total_loss = data_loss + loss_b + loss_ic + loss_pde + kl_weight * kl
    total_loss.backward()
    optimizer.step()
    if epoch % 500 == 0:
        print(f'Epoch {epoch}, Loss: {total_loss.item():.4f}')

# ---------- 验证与不确定性校准 ----------
# 测试网格
x_test = torch.linspace(0.1, 0.9, 20).view(-1, 1)
t_test = torch.linspace(0.1, 0.9, 20).view(-1, 1)
X, T = torch.meshgrid(x_test.squeeze(), t_test.squeeze(), indexing='ij')
x_flat = X.reshape(-1, 1)
t_flat = T.reshape(-1, 1)

# 预测（多次采样估计均值和方差）
n_mc = 50
pred_means = []
pred_stds = []
with torch.no_grad():
    for _ in range(n_mc):
        mean, std = model(x_flat, t_flat, n_samples=1)
        pred_means.append(mean[0])
        pred_stds.append(std[0])
pred_mean = torch.stack(pred_means).mean(dim=0)
pred_std = torch.stack(pred_means).std(dim=0)  # 认知不确定性

# 与精确解对比
u_true = u_exact(x_flat, t_flat)
error = (pred_mean - u_true).abs()
print(f'\n平均误差: {error.mean().item():.4f}')
print(f'预测标准差均值: {pred_std.mean().item():.4f}')

# 校准评估：计算Z-score（应接近标准正态）
z_score = (u_true - pred_mean) / (pred_std + 1e-6)
within_1sigma = (z_score.abs() < 1).float().mean().item()
within_2sigma = (z_score.abs() < 2).float().mean().item()
print(f'Z-score在1σ内比例: {within_1sigma:.3f} (理想≈0.683)')
print(f'Z-score在2σ内比例: {within_2sigma:.3f} (理想≈0.954)')

运行与验证

直接运行上述脚本即可完成训练与评估。训练约需1-2分钟（CPU）。判定标准：

• 训练损失持续下降并收敛；
• 平均误差（MAE）小于0.05（与噪声水平相当）；
• Z-score在1σ内的比例接近0.683，在2σ内接近0.954，说明不确定性校准良好。

关键说明

本实现采用均值网络计算PDE残差（忽略权重不确定性对残差的影响），这是B-PINN的常见简化处理。预测标准差由输出层权重方差传播得到，反映了认知不确定性。若需更精确的不确定性传播，可对残差也进行多次采样，但会显著增加计算量。KL权重需根据数据噪声调整，本案例取1e-4。