[API] Add chamfer loss (#871)

* add chamfer loss

* update development.md

* update example of chamfer loss docstring

* fix doctest for ChamferLoss

* support batch comupute for chamferloss

Author: HydrogenSulfate

README.md

@@ -220,9 +220,9 @@ PaddleScience 项目欢迎并依赖开发人员和开源社区中的用户，会
     旨在鼓励更多的开发者参与到飞桨科学计算社区的开源建设中，帮助社区修复 bug 或贡献 feature，加入开源、共建飞桨。了解编程基本知识的入门用户即可参与，活动进行中：
     [PaddleScience 快乐开源活动表单](https://github.com/PaddlePaddle/PaddleScience/issues/379)
 
-- 🔥第五期黑客松
+- 🔥第六期黑客松
 
-    面向全球开发者的深度学习领域编程活动，鼓励开发者了解与参与飞桨深度学习开源项目与文心大模型开发实践。活动进行中：[【PaddlePaddle Hackathon 5th】开源贡献个人挑战赛](https://github.com/PaddlePaddle/community/blob/master/hackathon/hackathon_5th/%E3%80%90PaddlePaddle%20Hackathon%205th%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#%E4%BB%BB%E5%8A%A1%E5%BC%80%E5%8F%91%E6%B5%81%E7%A8%8B%E4%B8%8E%E9%AA%8C%E6%94%B6%E6%A0%87%E5%87%86)
+    面向全球开发者的深度学习领域编程活动，鼓励开发者了解与参与飞桨深度学习开源项目与文心大模型开发实践。活动进行中：[【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/zh/api/loss/loss.md

@@ -11,6 +11,7 @@
         - L2RelLoss
         - MAELoss
         - MSELoss
+        - ChamferLoss
         - CausalMSELoss
         - MSELossWithL2Decay
         - IntegralLoss

docs/zh/development.md

@@ -116,15 +116,15 @@ model = ppsci.arch.MLP(("x", "y"), ("u", "v", "p"), 9, 50, "tanh")
 
         ``` py
         --8<--
-        ppsci/arch/mlp.py:86:151
+        ppsci/arch/mlp.py:139:279
         --8<--
         ```
 
     === "MLP.forward"
 
         ``` py
         --8<--
-        ppsci/arch/mlp.py:153:180
+        ppsci/arch/mlp.py:298:315
         --8<--
         ```
 

docs/zh/examples/allen_cahn.md

@@ -0,0 +1,229 @@
+# Allen-Cahn
+
+<!-- <a href="TODO" class="md-button md-button--primary" style>AI Studio快速体验</a> -->
+
+=== "模型训练命令"
+
+    ``` sh
+    python allen_cahn_default.py
+    ```
+
+=== "模型评估命令"
+
+    ``` sh
+    python allen_cahn_default.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/allen_cahn/allen_cahn_default_pretrained.pdparams
+    ```
+
+=== "模型导出命令"
+
+    ``` sh
+    python allen_cahn_default.py mode=export
+    ```
+
+=== "模型推理命令"
+
+    ``` sh
+    python allen_cahn_default.py mode=infer
+    ```
+
+| 预训练模型  | 指标 |
+|:--| :--|
+| [allen_cahn_default_pretrained.pdparams](TODO) | TODO |
+
+## 1. 背景简介
+
+Allen-Cahn 方程（有时也叫作模型方程或相场方程）是一种数学模型，通常用于描述两种不同相之间的界面演化。这个方程最早由Samuel Allen和John Cahn在1970年代提出，用以描述合金中相分离的过程。Allen-Cahn 方程是一种非线性偏微分方程，其一般形式可以写为：
+
+$$ \frac{\partial u}{\partial t} = \varepsilon^2 \Delta u - F'(u) $$
+
+这里：
+
+- $u(\mathbf{x},t)$ 是一个场变量，代表某个物理量，例如合金的组分浓度或者晶体中的有序参数。
+- $t$ 表示时间。
+- $\mathbf{x}$ 表示空间位置。
+- $\Delta$ 是Laplace算子，对应于空间变量的二阶偏导数（即 $\Delta u = \nabla^2 u$ ），用来描述空间扩散过程。
+- $\varepsilon$ 是一个正的小参数，它与相界面的宽度相关。
+- $F(u)$ 是一个双稳态势能函数，通常取为$F(u) = \frac{1}{4}(u^2-1)^2$，这使得 $F'(u) = u^3 - u$ 是其导数，这代表了非线性的反应项，负责驱动系统向稳定状态演化。
+
+这个方程中的 $F'(u)$ 项使得在 $u=1$ 和 $u=-1$ 附近有两个稳定的平衡态，这对应于不同的物理相。而 $\varepsilon^2 \Delta u$ 项则描述了相界面的曲率引起的扩散效应，这导致界面趋向于减小曲率。因此，Allen-Cahn 方程描述了由于相界面曲率和势能影响而发生的相变。
+
+在实际应用中，该方程还可能包含边界条件和初始条件，以便对特定问题进行数值模拟和分析。例如，在特定的物理问题中，可能会有 Neumann 边界条件（导数为零，表示无通量穿过边界）或 Dirichlet 边界条件（固定的边界值）。
+
+本案例解决以下 Allen-Cahn 方程：
+
+$$
+\begin{aligned}
+    & u_t - 0.0001 u_{xx} + 5 u^3 - 5 u  = 0,\quad t \in [0, 1],\ x\in[-1, 1],\\
+    &u(x,0) = x^2 \cos(\pi x),\\
+    &u(t, -1) = u(t, 1),\\
+    &u_x(t, -1) = u_x(t, 1).
+\end{aligned}
+$$
+
+## 2. 问题定义
+
+根据上述方程，可知计算域为$[0, 1]\times [-1, 1]$，含有一个初始条件： $u(x,0) = x^2 \cos(\pi x)$，两个周期边界条件：$u(t, -1) = u(t, 1)$、$u_x(t, -1) = u_x(t, 1)$。
+
+## 3. 问题求解
+
+接下来开始讲解如何将问题一步一步地转化为 PaddleScience 代码，用深度学习的方法求解该问题。
+为了快速理解 PaddleScience，接下来仅对模型构建、方程构建、计算域构建等关键步骤进行阐述，而其余细节请参考 [API文档](../api/arch.md)。
+
+### 3.1 模型构建
+
+在 Allen-Cahn 问题中，每一个已知的坐标点 $(t, x)$ 都有对应的待求解的未知量 $(u)$，
+，在这里使用比较简单的 MLP(Multilayer Perceptron, 多层感知机) 来表示 $(t, x)$ 到 $(u)$ 的映射函数 $f: \mathbb{R}^2 \to \mathbb{R}^1$ ，即：
+
+$$
+u = f(t, x)
+$$
+
+上式中 $f$ 即为 MLP 模型本身，用 PaddleScience 代码表示如下
+
+``` py linenums="63"
+--8<--
+examples/allen_cahn/allen_cahn_default.py:63:64
+--8<--
+```
+
+为了在计算时，准确快速地访问具体变量的值，在这里指定网络模型的输入变量名是 `("t", "x")`，输出变量名是 `("u")`，这些命名与后续代码保持一致。
+
+接着通过指定 MLP 的层数、神经元个数，就实例化出了一个拥有 4 层隐藏神经元，每层神经元数为 256 的神经网络模型 `model`，使用 `tanh` 作为激活函数。
+
+``` yaml linenums="35"
+--8<--
+examples/allen_cahn/conf/allen_cahn_default.yaml:35:41
+--8<--
+```
+
+### 3.2 方程构建
+
+Allen-Cahn 微分方程可以用如下代码表示：
+
+``` py linenums="66"
+--8<--
+examples/allen_cahn/allen_cahn_default.py:66:67
+--8<--
+```
+
+### 3.3 计算域构建
+
+本问题的计算域为 $[0, 1]\times [-1, 1]$，其中用于训练的数据已提前生成，保存在 `./dataset/allen_cahn.mat` 中，读取并生成计算域内的离散点。
+
+``` py linenums="69"
+--8<--
+examples/allen_cahn/allen_cahn_default.py:69:81
+--8<--
+```
+
+### 3.4 约束构建
+
+#### 3.4.1 内部点约束
+
+以作用在内部点上的 `SupervisedConstraint` 为例，代码如下：
+
+``` py linenums="94"
+--8<--
+examples/allen_cahn/allen_cahn_default.py:94:110
+--8<--
+```
+
+`SupervisedConstraint` 的第一个参数是用于训练的数据配置，由于我们使用实时随机生成的数据，而不是固定数据点，因此填入自定义的输入数据/标签生成函数；
+
+第二个参数是方程表达式，因此传入 Allen-Cahn 的方程对象；
+
+第三个参数是损失函数，此处选用 `CausalMSELoss` 函数，其会根据 `causal` 和 `tol` 参数，对不同的时间窗口进行重新加权， 能更好地优化瞬态问题；
+
+第四个参数是约束条件的名字，需要给每一个约束条件命名，方便后续对其索引。此处命名为 "PDE" 即可。
+
+#### 3.4.2 周期边界约束
+
+此处我们采用 hard-constraint 的方式，在神经网络模型中，对输入数据使用cos、sin等周期函数进行周期化，从而让$u_{\theta}$在数学上直接满足方程的周期性质。
+根据方程可得函数$u(t, x)$在$x$轴上的周期为2，因此将该周期设置到模型配置里即可。
+
+``` yaml linenums="35"
+--8<--
+examples/allen_cahn/conf/allen_cahn_default.yaml:35:43
+--8<--
+```
+
+#### 3.4.3 初值约束
+
+第三个约束条件是初值约束，代码如下：
+
+``` py linenums="112"
+--8<--
+examples/allen_cahn/allen_cahn_default.py:112:125
+--8<--
+```
+
+在微分方程约束、初值约束构建完毕之后，以刚才的命名为关键字，封装到一个字典中，方便后续访问。
+
+``` py linenums="126"
+--8<--
+examples/allen_cahn/allen_cahn_default.py:126:130
+--8<--
+```
+
+### 3.5 超参数设定
+
+接下来需要指定训练轮数和学习率，此处按实验经验，使用 200 轮训练轮数，0.001 的初始学习率。
+
+``` yaml linenums="51"
+--8<--
+examples/allen_cahn/conf/allen_cahn_default.yaml:51:73
+--8<--
+```
+
+### 3.6 优化器构建
+
+训练过程会调用优化器来更新模型参数，此处选择较为常用的 `Adam` 优化器，并配合使用机器学习中常用的 ExponentialDecay 学习率调整策略。
+
+``` py linenums="132"
+--8<--
+examples/allen_cahn/allen_cahn_default.py:132:136
+--8<--
+```
+
+### 3.7 评估器构建
+
+在训练过程中通常会按一定轮数间隔，用验证集（测试集）评估当前模型的训练情况，因此使用 `ppsci.validate.SupervisedValidator` 构建评估器。
+
+``` py linenums="138"
+--8<--
+examples/allen_cahn/allen_cahn_default.py:138:156
+--8<--
+```
+
+### 3.9 模型训练、评估与可视化
+
+完成上述设置之后，只需要将上述实例化的对象按顺序传递给 `ppsci.solver.Solver`，然后启动训练、评估、可视化。
+
+``` py linenums="158"
+--8<--
+examples/allen_cahn/allen_cahn_default.py:158:194
+--8<--
+```
+
+## 4. 完整代码
+
+``` py linenums="1" title="allen_cahn_default.py"
+--8<--
+examples/allen_cahn/allen_cahn_default.py
+--8<--
+```
+
+## 5. 结果展示
+
+在计算域上均匀采样出 $201\times501$ 个点，其预测结果和解析解如下图所示。
+
+<figure markdown>
+  ![allen_cahn_default.jpg](https://paddle-org.bj.bcebos.com/paddlescience/docs/AllenCahn/allen_cahn_default.png){ loading=lazy }
+  <figcaption> 左侧为 PaddleScience 预测结果，中间为解析解结果，右侧为两者的差值</figcaption>
+</figure>
+
+可以看到对于函数$u(t, x)$，模型的预测结果和解析解的结果基本一致。
+
+## 6. 参考资料
+
+- [Allen-Cahn equation](https://github.com/PredictiveIntelligenceLab/jaxpi/blob/main/examples/allen_cahn/README.md)

ppsci/loss/__init__.py

@@ -15,6 +15,7 @@
 import copy
 
 from ppsci.loss.base import Loss
+from ppsci.loss.chamfer import ChamferLoss
 from ppsci.loss.func import FunctionalLoss
 from ppsci.loss.integral import IntegralLoss
 from ppsci.loss.kl import KLLoss
@@ -40,6 +41,7 @@
     "PeriodicL2Loss",
     "MAELoss",
     "CausalMSELoss",
+    "ChamferLoss",
     "MSELoss",
     "MSELossWithL2Decay",
     "PeriodicMSELoss",

ppsci/loss/chamfer.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Dict
+from typing import Optional
+from typing import Union
+
+import paddle
+
+from ppsci.loss import base
+
+
+class ChamferLoss(base.Loss):
+    r"""Class for Chamfe distance loss.
+
+    $$
+    L = \dfrac{1}{S_1} \sum_{x \in S_1} \min_{y \in S_2} \Vert x - y \Vert_2^2 + \dfrac{1}{S_2} \sum_{y \in S_2} \min_{x \in S_1} \Vert y - x \Vert_2^2
+    $$
+
+    $$
+    \text{where } S_1 \text{ and } S_2 \text{ is the coordinate matrix of two point clouds}.
+    $$
+
+    Args:
+        weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None.
+
+    Examples:
+        >>> import paddle
+        >>> from ppsci.loss import ChamferLoss
+        >>> _ = paddle.seed(42)
+        >>> batch_point_cloud1 = paddle.rand([2, 100, 3])
+        >>> batch_point_cloud2 = paddle.rand([2, 50, 3])
+        >>> output_dict = {"s1": batch_point_cloud1}
+        >>> label_dict  = {"s1": batch_point_cloud2}
+        >>> weight = {"s1": 0.8}
+        >>> loss = ChamferLoss(weight=weight)
+        >>> result = loss(output_dict, label_dict)
+        >>> print(result)
+        Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True,
+               0.04415882)
+    """
+
+    def __init__(
+        self,
+        weight: Optional[Union[float, Dict[str, float]]] = None,
+    ):
+        super().__init__("mean", weight)
+
+    def forward(self, output_dict, label_dict, weight_dict=None):
+        losses = 0.0
+        for key in label_dict:
+            s1 = output_dict[key]
+            s2 = label_dict[key]
+            N1, N2 = s1.shape[1], s2.shape[1]
+
+            # [B, N1, N2, 3]
+            s1_expand = paddle.expand(s1.reshape([-1, N1, 1, 3]), shape=[-1, N1, N2, 3])
+            # [B, N1, N2, 3]
+            s2_expand = paddle.expand(s2.reshape([-1, 1, N2, 3]), shape=[-1, N1, N2, 3])
+
+            dis = ((s1_expand - s2_expand) ** 2).sum(axis=3)  # [B, N1, N2]
+            loss_s12 = dis.min(axis=2)  # [B, N1]
+            loss_s21 = dis.min(axis=1)  # [B, N2]
+            loss = loss_s12.mean() + loss_s21.mean()
+
+            if weight_dict and key in weight_dict:
+                loss *= weight_dict[key]
+
+            if isinstance(self.weight, (float, int)):
+                loss *= self.weight
+            elif isinstance(self.weight, dict) and key in self.weight:
+                loss *= self.weight[key]
+
+            losses += loss
+        return losses

ppsci/loss/mse.py

@@ -110,7 +110,7 @@ class CausalMSELoss(base.Loss):
     where $w_i=\exp (-\epsilon \displaystyle\sum_{k=1}^{i-1} \mathcal{L}_r^k), i=2,3, \ldots, M.$
 
     Args:
-        n_chunks (int): Number of time windows split.
+        n_chunks (int): $M$, Number of split time windows.
         reduction (Literal["mean", "sum"], optional): Reduction method. Defaults to "mean".
         weight (Optional[Union[float, Dict[str, float]]]): Weight for loss. Defaults to None.
         tol (float, optional): Causal tolerance, i.e. $\epsilon$ in paper. Defaults to 1.0.