ENGLISH | 简体中文
偏微分方程（PDE）与众多物理现象和工程应用紧密相联，涵盖机翼设计、电磁场模拟、应力分析等多个领域。 在这些实际应用中，PDE 的求解往往需要反复进行。 传统的 PDE 求解算法尽管十分精确，但往往会消耗大量的计算资源和时间。 近年提出的基于深度学习的神经算子（neural operator）方法极大地提高了 PDE 的求解速度，但难以泛化到新的 PDE 形式，且往往遇到训练代价高和数据量有限的问题。
针对以上问题，我们开发了 PDEformer 系列模型。 这是一种可以直接处理几乎所有 PDE 形式的端到端解预测模型，无需对不同 PDE 进行定制化的架构设计与训练，从而明显降低了模型部署成本，提升了求解效率。 针对一维 PDE 的 PDEformer-1 模型已于之前开源。 当前针对二维 PDE 的 PDEformer-2 模型使用约 40TB 的数据集进行预训练，能够对具有不同方程定义域、边界条件、变量个数、时间依赖情况的二维方程进行直接推理，快速获得任意时空位置的预测解。 此外，作为正问题解算子的可微分代理模型，PDEformer-2 还可以用于求解各类反问题，基于有噪声的时空散点观测，估计方程中的标量系数、源项场或波速场。 这为模型支持包括流体、电磁等领域的众多物理现象、工程应用的研究打下了良好的基础。
我们考虑定义在
其中
如图所示，PDEformer-2 先将方程形式表示为一个计算图，利用标量、函数编码器将方程涉及的数值信息嵌入到计算图的节点特征当中。 接下来，PDEformer-2 通过 graph Transformer 编码这一计算图，并使用隐式神经表征（INR）对所得的隐向量进行解码，获得 PDE 的各个解分量在特定时空坐标下的预测值。 关于模型工作原理的更详细解释可以在 PDEformer-1 的介绍中找到。
对于二维方程中可能出现的复杂区域形状和边界位置，PDEformer-2 将其表示为符号距离函数（SDF），并使用函数编码器将这部分信息嵌入到计算图当中。 下图所示的例子展示了使用计算图表示方形区域上 Dirichlet 边界条件的方式：
首先确保 MindSpore 已成功安装，见 安装教程。 其他依赖库可以通过如下命令安装：
pip3 install -r pip-requirements.txt由于本仓库文件较多，以下仅列出主要文件。 完整的文件目录结构请参见 docs/FILE_TREE_CN.md。
./
│ dynamic_dataset_manager.py # 数据集动态缓冲区管理程序，在大数据量（超出本地磁盘容量）预训练开始前启动
│ inverse_function.py # 反问题代码，估计 PDE 中的函数（源项、波方程速度场）
│ inverse_scalar.py # 反问题代码，估计 PDE 中的标量（方程系数）
│ PDEformer_inference.ipynb # 英文版模型预测交互式 notebook
│ PDEformer_inference_CN.ipynb # 中文版模型预测交互式 notebook
│ pip-requirements.txt # Python 依赖库
│ preprocess_data.py # 预处理数据、生成计算图，并将结果保存到新建的辅助数据文件中
│ README.md # 英文版说明文档
│ README_CN.md # 中文版说明文档
│ train.py # 模型训练代码
├─configs # 配置文件目录
│ ...
├─docs # 附加说明文档
│ ...
├─scripts # 用于启动训练、微调、反问题求解的 shell 脚本
│ ...
└─src # 基础代码目录
...
我们在 configs/inference 文件夹下提供了不同参数量的 PDEformer-2 模型的配置文件。 具体如下表所示：
| 模型 | 参数量 | 配置文件 | 预训练权重文件 |
|---|---|---|---|
| PDEformer2-S | 27.75M | configs/inference/model-S.yaml | model-S.ckpt |
| PDEformer2-M | 71.07M | configs/inference/model-M.yaml | model-M.ckpt |
| PDEformer2-L | 82.65M | configs/inference/model-L.yaml | model-L.ckpt |
下面的示例代码展示了如何使用 PDEformer-2 预测给定 PDE 的解，以非线性守恒律方程 model-M.ckpt，
并将 configs/inference/model-M.yaml 中 model.load_ckpt 参数的值改为相应的权重文件路径。
import numpy as np
from mindspore import context
from src import load_config, get_model, PDENodesCollector
from src.inference import infer_plot_2d, x_fenc, y_fenc
# 基本设定
context.set_context(mode=context.PYNATIVE_MODE, device_target="CPU")
config, _ = load_config("configs/inference/model-M.yaml")
model = get_model(config)
# 指定待求解的 PDE
pde = PDENodesCollector()
u = pde.new_uf()
u_ic = np.sin(2 * np.pi * x_fenc) * np.cos(4 * np.pi * y_fenc)
pde.set_ic(u, u_ic, x=x_fenc, y=y_fenc)
pde.sum_eq0(pde.dt(u), pde.dx(pde.square(u)), pde.dy(-0.3 * u))
# 使用 PDEformer 预测解（空间分辨率 32）并绘图
pde_dag = pde.gen_dag(config)
x_plot, y_plot = np.meshgrid(np.linspace(0, 1, 32), np.linspace(0, 1, 32), indexing="ij")
u_pred = infer_plot_2d(model, pde_dag, x_plot, y_plot)更多的示例见 PDEformer_inference_CN.ipynb。
PDEformer-2 通过预训练掌握了多种不同类型 PDE 的正问题求解能力。 这些多样 PDE 求解任务背后的共享知识为 PDEformer-2 快速适应新方程的求解打下了良好的基础。 我们可以使用一些标准数据集对预训练后的 PDEformer-2 模型进行微调，以评估它适配新场景、新任务的能力。 这里以 PDEBench 数据集中的 2D 浅水波（径向溃坝）数据为例，仅使用单个样本微调模型，并测试模型在 100 个测试样本上的预测精度。 微调过程依照如下几个步骤进行：
- (1) 从 Gitee AI 下载经过预训练的 PDEformer-2 权重
model-M.ckpt； - (2) 下载 PDEBench 数据集中的 2D_rdb_NA_NA.h5；
- (3) 调整配置文件 configs/finetune/pdebench-swe-rdb_model-M.yaml，指定下载所得的的模型权重文件路径（
model.ckpt参数）与数据集所在的文件夹（data.path参数）： 在这个配置文件中，我们需要指定数据集的文件路径和文件名，以及预训练好的模型权重路径 ：
# ...
model:
# ...
load_ckpt: path/to/your/downloaded/model-M.ckpt # 预训练模型权重路径
data:
path: ../data_download # 数据集所在目录
num_samples_per_file:
train: 1 # 训练数据集样本数
test: 100 # 测试数据集样本数
# ...
single_pde:
param_name: rdb # 方程数据集类型（有时还有参数名），这里为径向溃坝（RDB）数据
train: [1] # 方程参数值。当前所用的 RDB 数据集不包含参数，因此随意设置即可
# ...- (4) 完成配置文件的修改后，通过运行如下命令来启动单机单卡的微调任务：
config_path=configs/finetune/pdebench-swe-rdb_model-M.yaml
python train.py -c $config_path --no_distributed --device_id 0GitHub id: TODO
email: TODO