XPINN实战：如何用Python实现高维偏微分方程的高效求解（附代码）

XPINN实战：用Python攻克高维偏微分方程的工程化方案

在科学计算和工程仿真领域，高维偏微分方程的求解一直是个令人头疼的难题。传统的数值方法，比如有限元或有限差分，一旦维度升高，计算量就会呈指数级爆炸，网格生成也变得异常复杂。对于从事流体力学、金融衍生品定价或多物理场耦合仿真的工程师和研究人员来说，这常常意味着需要在计算精度、时间和硬件成本之间做出痛苦的权衡。

最近几年，基于物理信息的神经网络异军突起，为这个问题带来了全新的思路。它不再依赖离散网格，而是用神经网络直接学习控制方程的解。然而，标准的PINN在处理复杂几何或高维问题时，也会遇到训练困难、收敛缓慢的瓶颈。这时，扩展物理信息神经网络的出现，就像一把精准的手术刀，通过域分解的策略，将一个大问题拆解成多个可以并行处理的小问题。这不仅仅是理论上的优雅，更带来了实实在在的效率提升和灵活性。今天，我们就抛开复杂的公式推导，直接从代码和工程实现的角度，看看如何用Python和主流的深度学习框架，将XPINN应用到你的实际问题中，特别是那些具有不规则计算域的高维PDE。

1. 为什么是XPINN？从PINN的瓶颈说起

在深入代码之前，我们得先搞清楚XPINN要解决的核心痛点。传统的PINN用一个单一的神经网络去逼近整个求解域上的解。它的损失函数通常由三部分组成：控制方程的残差、边界条件以及可能的初始条件。这种方法概念简洁，但在实践中，尤其是面对以下场景时，会显得力不从心：

• 高维问题：网络的输入维度就是PDE的自变量维度。当维度增加时，网络需要学习的函数空间急剧膨胀，导致需要更多的参数和训练数据（残差点），训练变得极其不稳定。
• 复杂几何：对于不规则的计算域，均匀采样可能效率低下，很多采样点落在域外；而非均匀采样又增加了复杂性。更重要的是，在区域内部物理性质突变或存在内部界面时，单一网络很难精确捕捉这些局部特征。
• 并行化困难：虽然神经网络的前向传播和反向传播可以GPU加速，但整个优化问题本身是一个整体，难以进行粗粒度的任务并行。

cPINN迈出了第一步，它通过区域分解，在每个子域上使用独立的子网络。这带来了模型容量增加和潜在并行化的好处。但cPINN通常要求子域界面上的通量守恒，这对于非守恒型方程或复杂界面来说，条件过于严格，实现起来很麻烦。

注意：这里的“通量”需要根据具体PDE定义，对于泊松方程是梯度的法向分量，对于NS方程则涉及应力和速度。

而XPINN的关键创新在于其灵活且通用的界面条件。它不再强制要求法向通量连续，而是采用了更“宽松”且普适的条件，例如：

1. 解的平均值连续：在界面上，相邻子网络预测的解应该接近它们的平均值。
2. 残差连续性：在界面上，不同子网络满足控制方程的程度（即残差）应该接近。

这种设计带来了几个直接优势：

• 几何适应性极强：你可以几乎任意地划分计算域，无论是凸的、非凸的，还是像迷宫一样复杂的形状。划分的依据可以是计算资源的分布，也可以是物理场本身的特点（例如，在物理量梯度大的区域划分更密的子域）。
• 方程普适性：不再局限于守恒型方程，任何类型的PDE理论上都可以应用。
• 实现更简单：无需计算界面法向，降低了代码复杂度。

下面的表格对比了三种方法的典型特点：

理解了这些，我们就知道，当你的问题维度很高、计算域形状怪异，或者你希望充分利用多卡或多机资源时，XPINN是一个非常有吸引力的选项。

2. 构建XPINN求解器：一个模块化的Python蓝图

理论说得再多，不如一行代码。我们设计一个模块化的XPINN求解器，它应该包含以下几个核心部分：

1. 子域定义模块：负责生成或描述每个子域的几何范围。
2. 子网络模块：定义每个子域对应的独立神经网络。
3. 采样模块：在每个子域内部、边界以及界面上生成训练点。
4. 损失函数模块：计算每个子网络的总损失，包含数据损失、PDE残差损失和界面耦合损失。
5. 训练循环模块：协调多个子网络的优化过程，支持顺序或并行训练。

让我们用一个具体的例子来贯穿始终：求解一个二维的带源项泊松方程，但我们将计算域划分为两个不规则且重叠的L形区域来模拟复杂几何。方程如下： [ -\nabla^2 u = f(x, y), \quad (x, y) \in \Omega ] 边界条件设为狄利克雷条件，我们已知其解析解为 ( u(x,y) = \sin(\pi x) \sin(\pi y) )，据此可推导出源项 ( f(x,y) )。

2.1 定义子域与网络架构

首先，我们定义两个子域Omega1和Omega2。为了简单演示，我们用函数来判断点是否在域内。

import torch import torch.nn as nn import numpy as np # 假设我们的全局计算域是 [0,1]x[0,1] # 定义子域1: 一个L形区域，例如 [0, 0.7]x[0,1] 并集 [0,1]x[0, 0.7] def in_omega1(x, y): return ((x <= 0.7) & (y <= 1.0)) | ((x <= 1.0) & (y <= 0.7)) # 定义子域2: 另一个L形区域，与子域1部分重叠，例如 [0.3, 1]x[0,1] 并集 [0,1]x[0.3, 1] def in_omega2(x, y): return ((x >= 0.3) & (y <= 1.0)) | ((x <= 1.0) & (y >= 0.3)) # 定义每个子域独立的神经网络 class SubPINN(nn.Module): def __init__(self, layers=[2, 50, 50, 50, 1], activation=nn.Tanh): super(SubPINN, self).__init__() self.net = nn.Sequential() for i in range(len(layers)-2): self.net.add_module(f'linear_{i}', nn.Linear(layers[i], layers[i+1])) self.net.add_module(f'act_{i}', activation()) self.net.add_module('output', nn.Linear(layers[-2], layers[-1])) def forward(self, x): # x: [batch_size, 2] (x, y 坐标) return self.net(x) # 创建两个子网络 pinn1 = SubPINN() pinn2 = SubPINN()

这里每个子网络结构相同，但你可以根据子域内解的复杂度，为它们设计不同的深度、宽度或激活函数，这是XPINN的一个灵活之处。

2.2 智能采样策略：内部点、边界点与界面点

采样是PINN类方法成功的关键。对于XPINN，我们需要为每个子域采集三类点：

• 残差点：在子域内部随机采样，用于计算PDE残差损失。
• 边界点：在子域的真实物理边界上采样，用于施加边界条件。
• 界面点：在子域与相邻子域的重叠区域（即我们定义的公共界面区域）采样，用于施加界面耦合条件。

def sample_points(domain_func, n_residual, n_boundary, n_interface, global_domain_bounds): """ 为一个子域采样点。 domain_func: 判断点是否在子域内的函数 n_*: 各类点的数量 global_domain_bounds: 全局域的边界，用于初步采样 """ # 采样残差点：在全局域采样，然后过滤 x_res = torch.rand(n_residual, 2) * (global_domain_bounds[1] - global_domain_bounds[0]) + global_domain_bounds[0] # 过滤出在子域内的点 mask_res = domain_func(x_res[:, 0], x_res[:, 1]) x_res = x_res[mask_res][:n_residual] # 确保数量，可能需要循环采样 # 采样边界点：这里简化处理，在实际复杂几何中需要更精确的边界采样方法 # 例如，可以在全局边界上采样再过滤，或使用几何库生成边界点 x_bnd = torch.rand(n_boundary, 2) * (global_domain_bounds[1] - global_domain_bounds[0]) + global_domain_bounds[0] # ... 复杂的边界点生成逻辑（此处省略，依赖于具体几何） # 假设我们有一个函数能生成准确的边界点 x_bnd # 采样界面点：在全局域采样，然后过滤出同时在当前子域和至少一个相邻子域内的点 # 这里为了演示，我们假设界面区域是预定义的（例如两个L形的重叠矩形区域） interface_bounds = [[0.3, 0], [0.7, 0.7]] # 示例重叠区域 x_int = torch.rand(n_interface, 2) x_int[:, 0] = x_int[:, 0] * (interface_bounds[1][0] - interface_bounds[0][0]) + interface_bounds[0][0] x_int[:, 1] = x_int[:, 1] * (interface_bounds[1][1] - interface_bounds[0][1]) + interface_bounds[0][1] return x_res.requires_grad_(True), x_bnd.requires_grad_(True), x_int.requires_grad_(True) # 为两个子域采样 x_res1, x_bnd1, x_int1 = sample_points(in_omega1, 1000, 200, 300, [[0,0], [1,1]]) x_res2, x_bnd2, x_int2 = sample_points(in_omega2, 1000, 200, 300, [[0,0], [1,1]])

提示：对于极其复杂的几何，建议使用专业的网格生成库（如gmsh）来生成边界点和内部点，然后将坐标导入PyTorch。界面点的生成也需要根据子域划分的几何信息来精确计算。

2.3 核心：实现XPINN的复合损失函数

损失函数是XPINN的灵魂。对于每个子网络q，其总损失 ( \mathcal{J}_q ) 可以写作：

[ \mathcal{J}q = w{data} \cdot \text{MSE}{data} + w{pde} \cdot \text{MSE}{pde} + w{interface} \cdot \text{MSE}_{interface} ]

其中，界面损失 ( \text{MSE}_{interface} ) 通常包含解的平均连续损失和残差连续损失。

def compute_loss(pinn, x_res, x_bnd, x_int, pinn_neighbor, x_int_neighbor, source_func, exact_solution_func): """ 计算一个子网络的损失。 pinn: 当前子网络 x_res, x_bnd, x_int: 当前子域的残差点、边界点、界面点 pinn_neighbor: 相邻子网络（在界面的另一侧） x_int_neighbor: 相邻子网络在相同界面位置的坐标点（应与x_int一一对应或可插值） """ losses = {} # 1. PDE残差损失 (在内部残差点上) u_res = pinn(x_res) # 计算二阶导数 (Laplacian) grad_u = torch.autograd.grad(u_res, x_res, grad_outputs=torch.ones_like(u_res), create_graph=True)[0] grad_u_x = grad_u[:, 0:1] grad_u_y = grad_u[:, 1:2] grad2_u_x = torch.autograd.grad(grad_u_x, x_res, grad_outputs=torch.ones_like(grad_u_x), create_graph=True)[0][:, 0:1] grad2_u_y = torch.autograd.grad(grad_u_y, x_res, grad_outputs=torch.ones_like(grad_u_y), create_graph=True)[0][:, 1:2] laplacian_u = grad2_u_x + grad2_u_y f_res = source_func(x_res) # 已知源项 residual = -laplacian_u - f_res losses['loss_pde'] = torch.mean(residual**2) # 2. 边界条件损失 (在边界点上) u_bnd = pinn(x_bnd) u_bnd_exact = exact_solution_func(x_bnd) losses['loss_bnd'] = torch.mean((u_bnd - u_bnd_exact)**2) # 3. 界面耦合损失 (在界面点上) u_int = pinn(x_int) u_int_neighbor = pinn_neighbor(x_int_neighbor) # 假设界面点坐标一致 # 3a. 解的平均连续损失 u_avg = (u_int + u_int_neighbor) / 2.0 losses['loss_interface_avg'] = torch.mean((u_int - u_avg)**2) # 3b. 残差连续损失 (可选，但通常有助于稳定) # 计算当前网络在界面点的残差 u_int.requires_grad_(True) grad_u_int = torch.autograd.grad(u_int, x_int, grad_outputs=torch.ones_like(u_int), create_graph=True)[0] grad_u_int_x = grad_u_int[:, 0:1] grad_u_int_y = grad_u_int[:, 1:2] grad2_u_int_x = torch.autograd.grad(grad_u_int_x, x_int, grad_outputs=torch.ones_like(grad_u_int_x), create_graph=True)[0][:, 0:1] grad2_u_int_y = torch.autograd.grad(grad_u_int_y, x_int, grad_outputs=torch.ones_like(grad_u_int_y), create_graph=True)[0][:, 1:2] laplacian_u_int = grad2_u_int_x + grad2_u_int_y f_int = source_func(x_int) residual_int = -laplacian_u_int - f_int # 计算相邻网络在相同界面点的残差 (需要重新计算梯度) u_int_neighbor.requires_grad_(True) grad_u_nb = torch.autograd.grad(u_int_neighbor, x_int_neighbor, grad_outputs=torch.ones_like(u_int_neighbor), create_graph=True)[0] grad_u_nb_x = grad_u_nb[:, 0:1] grad_u_nb_y = grad_u_nb[:, 1:2] grad2_u_nb_x = torch.autograd.grad(grad_u_nb_x, x_int_neighbor, grad_outputs=torch.ones_like(grad_u_nb_x), create_graph=True)[0][:, 0:1] grad2_u_nb_y = torch.autograd.grad(grad_u_nb_y, x_int_neighbor, grad_outputs=torch.ones_like(grad_u_nb_y), create_graph=True)[0][:, 1:2] laplacian_u_nb = grad2_u_nb_x + grad2_u_nb_y residual_int_neighbor = -laplacian_u_nb - f_int # 使用相同的源项 losses['loss_interface_res'] = torch.mean((residual_int - residual_int_neighbor)**2) # 加权总损失 w_pde, w_bnd, w_int_avg, w_int_res = 1.0, 10.0, 5.0, 2.0 # 损失权重需要调参 total_loss = (w_pde * losses['loss_pde'] + w_bnd * losses['loss_bnd'] + w_int_avg * losses['loss_interface_avg'] + w_int_res * losses['loss_interface_res']) losses['total'] = total_loss return losses

这段代码清晰地展示了XPINN损失函数的构成。权重w_*的选择至关重要，通常边界条件和界面条件的权重需要设置得比PDE残差权重大一些，以确保这些硬约束被优先满足。这是一个需要根据具体问题反复调试的超参数。

2.4 训练策略与并行化考量

有了损失函数，我们就可以训练了。XPINN的训练可以有两种模式：

• 交替训练：依次优化每个子网络，固定其他子网络的参数。这种方式实现简单，但可能收敛较慢。
• 联合训练：将所有子网络的参数放在一个优化器中同时优化。这种方式理论上更高效，但损失函数 landscape 更复杂，可能更难调优。

# 采用交替训练策略 optimizer1 = torch.optim.Adam(pinn1.parameters(), lr=1e-3) optimizer2 = torch.optim.Adam(pinn2.parameters(), lr=1e-3) epochs = 20000 for epoch in range(epochs): # 训练子网络1 (固定子网络2) pinn2.eval() optimizer1.zero_grad() loss_dict1 = compute_loss(pinn1, x_res1, x_bnd1, x_int1, pinn2, x_int1, source_func, exact_solution_func) loss1 = loss_dict1['total'] loss1.backward() optimizer1.step() # 训练子网络2 (固定子网络1) pinn1.eval() optimizer2.zero_grad() loss_dict2 = compute_loss(pinn2, x_res2, x_bnd2, x_int2, pinn1, x_int2, source_func, exact_solution_func) loss2 = loss_dict2['total'] loss2.backward() optimizer2.step() if epoch % 1000 == 0: print(f"Epoch {epoch}, Loss1: {loss1.item():.4e}, Loss2: {loss2.item():.4e}")

对于并行化，由于子网络之间的耦合仅通过界面损失项（在计算时需要对方网络的预测值），因此真正的数据并行（如DataParallel）并不直接适用。但我们可以采用模型并行的思路：将不同的子网络放在不同的GPU上，在计算界面损失时，通过GPU间的数据传输交换界面点的预测值。这需要更精细的工程实现，例如使用torch.distributed或multiprocessing模块。

3. 性能对比实验：XPINN vs. 标准PINN

空谈不如实验。我们在同一个二维泊松方程问题上，对比标准PINN和XPINN（2个子域）的表现。为了公平，我们控制总参数量大致相同（例如，标准PINN用一个 [2, 100, 100, 100, 1] 的网络，XPINN用两个 [2, 50, 50, 50, 1] 的网络）。

我们主要关注以下几个指标：

• 训练速度：达到相同验证集精度所需的epoch数或时间。
• 最终精度：在测试集上的相对L2误差。
• 内存占用：训练过程中的峰值GPU内存。
• 对复杂几何的适应性：在具有内部界面或奇异性的区域，解的精度。

在我的本地测试（RTX 3080）中，一个简化的对比结果如下：

注意：这个对比结果高度依赖于问题设置、网络结构、超参数和优化器。XPINN的优势在问题复杂度（维度和几何）增加时会更加明显。

实验代码的核心是比较两者的损失下降曲线和最终解的误差分布图。我们可以用matplotlib绘制误差云图，能直观地看到在传统PINN表现不佳的区域，XPINN是否有所改善。

# 评估代码片段示例 def evaluate_model(model, global_domain_points, exact_solution_func): model.eval() with torch.no_grad(): u_pred = model(global_domain_points) u_exact = exact_solution_func(global_domain_points) error = torch.abs(u_pred - u_exact) relative_l2_error = torch.norm(error) / torch.norm(u_exact) return relative_l2_error, error # 在整个计算域上生成密集的测试点 grid_x, grid_y = torch.meshgrid(torch.linspace(0,1,100), torch.linspace(0,1,100)) test_points = torch.stack([grid_x.flatten(), grid_y.flatten()], dim=1) # 对于XPINN，需要将测试点分配到各个子域，并用对应的子网络预测 def evaluate_xpinn(pinn1, pinn2, test_points, in_omega1, in_omega2): mask1 = in_omega1(test_points[:,0], test_points[:,1]) mask2 = in_omega2(test_points[:,0], test_points[:,1]) # 处理重叠区域：这里简单取平均，更复杂的策略可以按距离加权 u_pred_global = torch.zeros_like(test_points[:,0:1]) with torch.no_grad(): if mask1.any(): u_pred_global[mask1] += pinn1(test_points[mask1]) if mask2.any(): u_pred_global[mask2] += pinn2(test_points[mask2]) # 对于同时在两个子域的点，上面加了两次，需要平均 overlap_mask = mask1 & mask2 if overlap_mask.any(): u_pred_global[overlap_mask] /= 2.0 u_exact = exact_solution_func(test_points) error = torch.abs(u_pred_global - u_exact) relative_l2_error = torch.norm(error) / torch.norm(u_exact) return relative_l2_error, error, u_pred_global

4. 工程实践中的技巧与避坑指南

在实际项目中应用XPINN，有几个细节处理不好就容易翻车。

1. 子域划分的艺术划分不是随意的。一个好的划分应该：

• 平衡负载：让各个子域内的计算复杂度（如残差点数量）大致相当。
• 对齐物理特征：尽量让物理场平滑的区域在一个子域内，将间断面、奇异性或边界层放在子域界面附近，让界面条件去处理这些“麻烦”。
• 保证足够重叠：界面区域需要有足够的采样点，以确保耦合的有效性。重叠区域太薄会导致训练不稳定。

2. 损失权重的动态调整固定权重可能不是最优的。可以采用自适应权重策略，例如：

• 根据各个损失项的量级动态调整，使它们在训练初期保持同一数量级。
• 使用LR类似的思想，在训练过程中逐渐增加界面损失的权重，让网络先大致满足PDE和边界条件，再加强子域间的耦合。

# 一个简单的自适应权重策略示例 class AdaptiveWeight: def __init__(self, initial_weight, update_rate=0.01): self.weight = torch.tensor(initial_weight, requires_grad=False) self.update_rate = update_rate self.history = [] def update(self, current_loss_component): # 例如，让权重与该损失分量的历史平均值成反比 self.history.append(current_loss_component.item()) avg_loss = np.mean(self.history[-100:]) if len(self.history) > 100 else np.mean(self.history) self.weight = self.weight * (1 - self.update_rate) + (self.update_rate / (avg_loss + 1e-8))

3. 界面点的匹配与通信在并行计算中，子网络可能位于不同的进程或设备上。界面点x_int必须严格匹配，即两个子网络在计算界面损失时，针对的是空间中的同一组坐标点。这需要预先定义好全局的界面点集，或者设计一个插值例程，将一个子网络的界面预测值插值到另一个子网络的界面点上。

4. 处理逆问题与参数识别XPINN同样可以用于逆问题，即同时求解场和未知参数。此时，损失函数中需要增加对参数一致性的约束。在界面处，不仅解要连续，识别出的参数在界面两侧也应该一致。这需要在界面损失中加入一项参数差异的惩罚项，正如原始论文中MSE_lambda项所示。

5. 调试与可视化由于问题更复杂，调试至关重要。

• 监控每个子损失：分别记录loss_pde,loss_bnd,loss_interface_avg等，看是哪部分不收敛。
• 可视化子域解：训练过程中，定期绘制每个子网络在其子域上的预测结果，检查在界面处是否出现明显的跳跃或不连续。
• 验证界面条件：单独计算界面点上解的平均差和残差差，确保它们随着训练而减小。

最后，别忘了社区的力量。DeepXDE、SimNet等开源库已经提供了PINN的高级实现，虽然它们可能没有直接提供XPINN的封装，但其代码结构（如损失定义、采样、训练循环）是极好的参考。从修改这些成熟的代码库开始，比自己从头搭建要高效得多。

将XPINN从论文公式落地到可运行的Python代码，最大的收获不是调通了几个参数，而是对整个“分解-协调”思想有了具象的理解。它本质上是一种分而治之的策略，用多个小型、专用的神经网络去协作解决一个大型、复杂的问题。这种思路不仅可以用于PDE求解，对于其他高维函数逼近、多物理场耦合问题也有启发意义。在实际项目中，我通常会先用一个小规模的、几何规则的问题验证XPINN流程，确保界面耦合和训练循环正确无误，然后再逐步扩展到真实的高维复杂场景中。