[Example] Add spinn on helmholtz equation (#958)

* update spinn code(WIP)

* update code(l2err=0.089)

* update SPINN helmholtz3d

* update pretraind model url(TRAIN.nc=64)

* add spinn doc

* refine code and docs

* update

* update docs

* update supports information

* update code and doc

* refine arch code

Author: HydrogenSulfate

README.md

@@ -27,6 +27,7 @@ PaddleScience 是一个基于深度学习框架 PaddlePaddle 开发的科学计
 
 | 问题类型 | 案例名称 | 优化算法 | 模型类型 | 训练方式 | 数据集 | 参考资料 |
 |-----|---------|-----|---------|----|---------|---------|
+| 三维亥姆霍兹方程 | [SPINN(Helmholtz3D)](https://paddlescience-docs.readthedocs.io/zh/examples/spinn.md) | 机理驱动 | ModifiedMLP | 无监督学习 | - | [Paper](https://arxiv.org/pdf/2306.15969) |
 | 相场方程 | [Allen-Cahn](https://paddlescience-docs.readthedocs.io/zh/latest/zh/examples/allen_cahn) | 机理驱动 | MLP | 无监督学习 | [Data](https://paddle-org.bj.bcebos.com/paddlescience/datasets/AllenCahn/allen_cahn.mat) | [Paper](https://arxiv.org/pdf/2402.00326) |
 | 微分方程 | [拉普拉斯方程](https://paddlescience-docs.readthedocs.io/zh/latest/zh/examples/laplace2d) | 机理驱动 | MLP | 无监督学习 | -        | - |
 | 微分方程 | [伯格斯方程](https://paddlescience-docs.readthedocs.io/zh/latest/zh/examples/deephpms) | 机理驱动 | MLP | 无监督学习 | [Data](https://github.com/maziarraissi/DeepHPMs/tree/master/Data) | [Paper](https://arxiv.org/pdf/1801.06637.pdf) |
@@ -212,7 +213,7 @@ python -c "import paddle; paddle.utils.run_check()"
 
 ## 🎈其他领域支持
 
-除 PaddleScience 套件外，Paddle 框架还支持了 [DeepXDE](https://github.com/lululxvi/deepxde/tree/master?tab=readme-ov-file#deepxde) 的所有案例，分子动力学套件 [DeepMD-kit](https://github.com/deepmodeling/deepmd-kit/tree/paddle2?tab=readme-ov-file#deepmd-kitpaddlepaddle-backend) 部分案例和功能，以及正在适配中的 Modulus。
+除 PaddleScience 套件外，Paddle 框架还支持了 [Modulus-sym](https://github.com/PaddlePaddle/modulus-sym/tree/paddle?tab=readme-ov-file#modulus-symbolic-betapaddle-backend)、[DeepXDE](https://github.com/lululxvi/deepxde/tree/master?tab=readme-ov-file#deepxde) 的所有案例，分子动力学套件 [DeepMD-kit](https://github.com/deepmodeling/deepmd-kit/tree/paddle2?tab=readme-ov-file#deepmd-kitpaddlepaddle-backend) 部分案例和功能。
 
 <!-- --8<-- [start:support] -->
 ## 💬支持与建议
@@ -232,9 +233,9 @@ PaddleScience 项目欢迎并依赖开发人员和开源社区中的用户，会
     旨在鼓励更多的开发者参与到飞桨科学计算社区的开源建设中，帮助社区修复 bug 或贡献 feature，加入开源、共建飞桨。了解编程基本知识的入门用户即可参与，活动进行中：
     [PaddleScience 快乐开源活动表单](https://github.com/PaddlePaddle/PaddleScience/issues/379)
 
-- 🔥第六期黑客松
+<!-- - 🔥第六期黑客松
 
-    面向全球开发者的深度学习领域编程活动，鼓励开发者了解与参与飞桨深度学习开源项目与文心大模型开发实践。活动进行中：[【PaddlePaddle Hackathon 5th】开源贡献个人挑战赛](https://github.com/PaddlePaddle/community/blob/master/hackathon/hackathon_6th/%E3%80%90Hackathon%206th%E3%80%91%E5%BC%80%E6%BA%90%E8%B4%A1%E7%8C%AE%E4%B8%AA%E4%BA%BA%E6%8C%91%E6%88%98%E8%B5%9B%E7%A7%91%E5%AD%A6%E8%AE%A1%E7%AE%97%E4%BB%BB%E5%8A%A1%E5%90%88%E9%9B%86.md)
+    面向全球开发者的深度学习领域编程活动，鼓励开发者了解与参与飞桨深度学习开源项目与文心大模型开发实践。活动进行中：[【PaddlePaddle Hackathon 5th】开源贡献个人挑战赛](https://github.com/PaddlePaddle/community/blob/master/hackathon/hackathon_6th/%E3%80%90Hackathon%206th%E3%80%91%E5%BC%80%E6%BA%90%E8%B4%A1%E7%8C%AE%E4%B8%AA%E4%BA%BA%E6%8C%91%E6%88%98%E8%B5%9B%E7%A7%91%E5%AD%A6%E8%AE%A1%E7%AE%97%E4%BB%BB%E5%8A%A1%E5%90%88%E9%9B%86.md) -->
 <!-- --8<-- [end:contribution] -->
 
 <!-- --8<-- [start:collaboration] -->

docs/index.md

@@ -72,6 +72,7 @@
 
 | 问题类型 | 案例名称 | 优化算法 | 模型类型 | 训练方式 | 数据集 | 参考资料 |
 |-----|---------|-----|---------|----|---------|---------|
+| 三维亥姆霍兹方程 | [SPINN(Helmholtz3D)](./zh/examples/spinn.md) | 机理驱动 | ModifiedMLP | 无监督学习 | - | [Paper](https://arxiv.org/pdf/2306.15969) |
 | 相场方程 | [Allen-Cahn](./zh/examples/allen_cahn.md) | 机理驱动 | MLP | 无监督学习 | [Data](https://paddle-org.bj.bcebos.com/paddlescience/datasets/AllenCahn/allen_cahn.mat) | [Paper](https://arxiv.org/pdf/2402.00326) |
 | 微分方程 | [拉普拉斯方程](./zh/examples/laplace2d.md) | 机理驱动 | MLP | 无监督学习 | -        | - |
 | 微分方程 | [伯格斯方程](./zh/examples/deephpms.md) | 机理驱动 | MLP | 无监督学习 | [Data](https://github.com/maziarraissi/DeepHPMs/tree/master/Data) | [Paper](https://arxiv.org/pdf/1801.06637.pdf) |
@@ -181,16 +182,16 @@
             <br><span class="text-large">全量支持</span></br>
         </div>
     </a>
+    <a href="https://github.com/PaddlePaddle/modulus-sym/tree/paddle?tab=readme-ov-file#modulus-symbolic-betapaddle-backend">
+        <div class="card card-deepmd">
+            Modulus-sym
+            <br><span class="text-large">全量支持</span></br>
+        </div>
+    </a>
     <a href="https://github.com/deepmodeling/deepmd-kit/tree/paddle?tab=readme-ov-file#deepmd-kitpaddlepaddle-backend">
         <div class="card card-modulus">
             DeepMD
-            <br><span class="text-large">适配中</span></br>
-        </div>
-    </a>
-    <a href="https://github.com/PaddlePaddle/modulus-sym/tree/paddle?tab=readme-ov-file#modulus-symbolic-betapaddle-backend">
-        <div class="card card-deepmd">
-            Modulus
-            <br><span class="text-large">适配中</span></br>
+            <br><span class="text-large">部分适配</span></br>
         </div>
     </a>
 </div>

docs/zh/api/equation.md

@@ -9,6 +9,7 @@
         - Biharmonic
         - FractionalPoisson
         - HeatExchanger
+        - Helmholtz
         - Laplace
         - LinearElasticity
         - NavierStokes

docs/zh/examples/spinn.md

@@ -0,0 +1,205 @@
+# SPINN(helmholtz3d)
+
+<!-- <a href="https://aistudio.baidu.com/projectdetail/8219967" class="md-button md-button--primary" style>AI Studio快速体验</a> -->
+
+=== "模型训练命令"
+
+    ``` sh
+    python helmholtz3d.py
+    ```
+
+=== "模型评估命令"
+
+    ``` sh
+    python helmholtz3d.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/spinn/spinn_helmholtz3d_pretrained.pdparams
+    ```
+
+=== "模型导出命令"
+
+    ``` sh
+    python helmholtz3d.py mode=export
+    ```
+
+=== "模型推理命令"
+
+    ``` sh
+    python helmholtz3d.py mode=infer
+    ```
+
+| 预训练模型  | 指标 |
+|:--| :--|
+| [spinn_helmholtz3d.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/spinn/spinn_helmholtz3d_pretrained.pdparams) | l2_err: 0.0183 <br> rmse: 0.0064 |
+
+## 1. 背景简介
+
+Helmholtz方程是一个重要的偏微分方程，广泛应用于物理学和工程学中，特别是在波动理论和振动问题中。它以德国物理学家赫尔曼·冯·亥姆霍兹（Hermann von Helmholtz）的名字命名。Helmholtz方程的标准形式如下：
+
+$$
+\nabla^2 u + k^2 u = q
+$$
+
+这里：
+
+- $\nabla^2$ 是拉普拉斯算子（也称为拉普拉斯算符），在三维直角坐标系下，它的形式是：$\nabla^2 = \frac{\partial^2 }{\partial x^2} + \frac{\partial^2 }{\partial y^2} + \frac{\partial^2 }{\partial z^2}$
+- $u$ 是待求解的函数，通常表示物理量的幅度，如电磁场、声压或量子波函数等。
+- $k$ 是波数，定义为 $k = \frac{2\pi}{\lambda}$，其中 $\lambda$ 是波长。
+- $q$ 是源项，通常表示物理量与时间、空间导数之间的相互作用。
+
+本案例解决以下三维 Helmholtz 方程：
+
+$$
+\begin{aligned}
+  & \nabla^2 u + k^2 u = q, x \in \Omega \\
+  & u(x) = 0, x \in \partial \Omega \\
+\end{aligned}
+$$
+
+$$
+\begin{aligned}
+  & \text{source term } q = -(a_1 \pi)^2 u -(a_2 \pi)^2 u -(a_3 \pi)^2 u + k^2 u \\
+  & \text{where }k=1, a_1=4, a_2=4, a_3=3
+\end{aligned}
+$$
+
+## 2. 问题定义
+
+本问题的计算域在 $[-1, 1] ^3$ 一个单位正方体内，对于计算域内部点，要求满足上述 Helmholtz 方程，对于计算域边界点，要求满足 $u = 0$。
+
+## 3. 问题求解
+
+接下来开始讲解如何将问题一步一步地转化为 PaddleScience 代码，用深度学习的方法求解该问题。
+为了快速理解 PaddleScience，接下来仅对模型构建、方程构建、计算域构建等关键步骤进行阐述，而其余细节请参考 [API文档](../api/arch.md)。
+
+### 3.1 模型构建
+
+SPINN 的模型结构设计如下：
+
+![SPINN_structure](https://paddle-org.bj.bcebos.com/paddlescience/docs/spinn/spinn_structure.png)
+
+在 Helmholtz 问题中，每一个已知的坐标点 $(x, y, z)$ 都有对应的待求解的未知量 $u$（此处我们用 $u$代替），在这里使用 SPINN 来表示 $(x, y, z)$ 到 $(u)$ 的映射函数 $f: \mathbb{R}^3 \to \mathbb{R}^1$ ，即：
+
+$$
+u = m(x, y, z)
+$$
+
+上式中 $m$ 即为 SPINN 模型本身，用 PaddleScience 代码表示如下
+
+``` py linenums="99"
+--8<--
+examples/spinn/helmholtz3d.py:99:100
+--8<--
+```
+
+为了在计算时，准确快速地访问具体变量的值，在这里指定网络模型的输入变量名是 `("x", "y", "z")`，输出变量名是 `("u")`，这些命名与后续代码保持一致。
+
+接着通过指定 SPINN 的层数、神经元个数，就实例化出了一个拥有 4 层全连接层，每层全连接层的神经元个数为 64 ，每一个输出变量的隐层特征维度 `r` 为 32 的神经网络模型 `model`， 并且使用 `tanh` 作为激活函数。
+
+``` yaml linenums="38"
+--8<--
+examples/spinn/conf/helmholtz3d.yaml:38:45
+--8<--
+```
+
+### 3.2 方程构建
+
+Helmholtz 微分方程可以用如下代码表示：
+
+``` py linenums="102"
+--8<--
+examples/spinn/helmholtz3d.py:102:104
+--8<--
+```
+
+注：此处我们需要把 model 手动传递给 `equation["Helmholtz"]`，因为 `Helmholtz` 方程需要用到前向微分功能。
+
+### 3.3 约束构建
+
+#### 3.3.1 内部点约束
+
+以作用在内部点上的 `SupervisedConstraint` 为例，用于生成内部点训练数据的代码如下：
+
+``` py linenums="39"
+--8<--
+examples/spinn/helmholtz3d.py:39:83
+--8<--
+```
+
+用于构建内部点约束的代码如下：
+
+``` py linenums="106"
+--8<--
+examples/spinn/helmholtz3d.py:106:156
+--8<--
+```
+
+`SupervisedConstraint` 的第一个参数是用于训练的数据配置，由于我们使用实时随机生成的数据，而不是固定数据点，因此填入自定义的输入数据/标签生成函数；
+
+第二个参数是方程表达式，因此传入 Helmholtz 的方程对象；
+
+第三个参数是损失函数，此处选用 `MSELoss` 即可；
+
+第四个参数是约束条件的名字，需要给每一个约束条件命名，方便后续对其索引。此处命名为 "PDE" 即可。
+
+#### 3.3.3 边值约束
+
+第三个约束条件是边值约束，代码如下：
+
+``` py linenums="158"
+--8<--
+examples/spinn/helmholtz3d.py:158:190
+--8<--
+```
+
+### 3.4 超参数设定
+
+接下来需要指定训练轮数和学习率，此处按实验经验，使用 50 轮训练轮数，每轮迭代 1000 步，0.001 的初始学习率。
+
+``` yaml linenums="47"
+--8<--
+examples/spinn/conf/helmholtz3d.yaml:47:63
+--8<--
+```
+
+### 3.5 优化器构建
+
+训练过程会调用优化器来更新模型参数，此处选择较为常用的 `Adam` 优化器，并配合使用机器学习中常用的 ExponentialDecay 学习率调整策略。
+
+``` py linenums="192"
+--8<--
+examples/spinn/helmholtz3d.py:192:196
+--8<--
+```
+
+### 3.6 模型训练、评估与可视化
+
+完成上述设置之后，只需要将上述实例化的对象按顺序传递给 `ppsci.solver.Solver`，然后启动训练、评估、可视化。
+
+``` py linenums="198"
+--8<--
+examples/spinn/helmholtz3d.py:198:227
+--8<--
+```
+
+## 4. 完整代码
+
+``` py linenums="1" title="helmholtz3d.py"
+--8<--
+examples/spinn/helmholtz3d.py
+--8<--
+```
+
+## 5. 结果展示
+
+在计算域上均匀采样出 $100^3$ 个点，其预测结果和解析解如下图所示。
+
+<figure markdown>
+  ![spinn_helmholtz3d.jpg](https://paddle-org.bj.bcebos.com/paddlescience/docs/spinn/spinn_helmholtz3d.png){ loading=lazy }
+  <figcaption> 左侧为 PaddleScience 预测结果，右侧为解析解结果</figcaption>
+</figure>
+
+本问题使用模型预测的误差为 l2_err = 0.0183，rmse = 0.0064，与解析解误差较小，基本一致。
+
+## 6. 参考资料
+
+- [Separable Physics-Informed Neural Networks](https://arxiv.org/pdf/2306.15969)
+- [SPINN](https://github.com/stnamjef/SPINN?tab=readme-ov-file)

examples/spinn/conf/helmholtz3d.yaml

@@ -0,0 +1,88 @@
+defaults:
+  - ppsci_default
+  - TRAIN: train_default
+  - TRAIN/ema: ema_default
+  - TRAIN/swa: swa_default
+  - EVAL: eval_default
+  - INFER: infer_default
+  - hydra/job/config/override_dirname/exclude_keys: exclude_keys_default
+  - _self_
+
+hydra:
+  run:
+    # dynamic output directory according to running time and override name
+    dir: outputs_spinn_helmholtz3d/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname}
+  job:
+    name: ${mode} # name of logfile
+    chdir: false # keep current working directory unchanged
+  callbacks:
+    init_callback:
+      _target_: ppsci.utils.callbacks.InitCallback
+  sweep:
+    # output directory for multirun
+    dir: ${hydra.run.dir}
+    subdir: ./
+
+# general settings
+mode: train # running mode: train/eval
+seed: 111
+output_dir: ${hydra:run.dir}
+log_freq: 100
+
+# set working condition
+K: 1.0
+a1: 4
+a2: 4
+a3: 3
+
+# model settings
+MODEL:
+  input_keys: ["x", "y", "z"]
+  output_keys: ["u"]
+  num_layers: 4
+  hidden_size: 64
+  r: 32
+  activation: "tanh"
+
+# training settings
+TRAIN:
+  epochs: 50
+  iters_per_epoch: 1000
+  save_freq: 10
+  eval_during_train: false
+  eval_freq: 5
+  lr_scheduler:
+    epochs: ${TRAIN.epochs}
+    iters_per_epoch: ${TRAIN.iters_per_epoch}
+    learning_rate: 1.0e-3
+    gamma: 0.9
+    decay_steps: 1000
+    by_epoch: false
+  nc: 64
+  pretrained_model_path: null
+  checkpoint_path: null
+
+# evaluation settings
+EVAL:
+  pretrained_model_path: null
+  eval_with_no_grad: true
+  batch_size: 1024
+  nc: 100
+
+# inference settings
+INFER:
+  pretrained_model_path: https://paddle-org.bj.bcebos.com/paddlescience/models/spinn/spinn_helmholtz3d_pretrained.pdparams
+  export_path: ./inference/spinn_helmholtz3d
+  pdmodel_path: ${INFER.export_path}.pdmodel
+  pdiparams_path: ${INFER.export_path}.pdiparams
+  onnx_path: ${INFER.export_path}.onnx
+  device: gpu
+  engine: native
+  precision: fp32
+  ir_optim: true
+  min_subgraph_size: 5
+  gpu_mem: 2000
+  gpu_id: 0
+  max_batch_size: 1024
+  num_cpu_threads: 10
+  batch_size: 1024