support dim=1/2/3 for SPINN and helmholtz PDE (#1075)

Author: HydrogenSulfate

examples/spinn/helmholtz3d.py

@@ -100,8 +100,7 @@ def train(cfg: DictConfig):
     model = ppsci.arch.SPINN(**cfg.MODEL)
 
     # set equation
-    equation = {"Helmholtz": ppsci.equation.Helmholtz(3, 1.0)}
-    equation["Helmholtz"].model = model  # set model to equation for hvp
+    equation = {"Helmholtz": ppsci.equation.Helmholtz(3, 1.0, model)}
 
     # set constraint
     class InteriorDataGenerator:

ppsci/arch/spinn.py

@@ -137,10 +137,10 @@ def _tensor_contraction(self, x: paddle.Tensor, y: paddle.Tensor) -> paddle.Tens
 
         return out
 
-    def forward_tensor(self, x, y, z) -> List[paddle.Tensor]:
+    def forward_tensor(self, *xs) -> List[paddle.Tensor]:
         # forward each dim branch
         feature_f = []
-        for i, input_var in enumerate((x, y, z)):
+        for i, input_var in enumerate(xs):
             input_i = {self.input_keys[i]: input_var}
             output_f_i = self.branch_nets[i](input_i)
             feature_f.append(output_f_i["f"])  # [B, r*output_dim]

ppsci/equation/pde/helmholtz.py

@@ -66,28 +66,50 @@ def __init__(
         self,
         dim: int,
         k: float,
+        model: paddle.nn.Layer,
         detach_keys: Optional[Tuple[str, ...]] = None,
     ):
         super().__init__()
         self.dim = dim
         self.k = k
         self.detach_keys = detach_keys
 
-        self.model: paddle.nn.Layer
+        invars = self.create_symbols("x y z")[:dim]
 
-        def helmholtz(data_dict: Dict[str, "paddle.Tensor"]):
-            x, y, z = (
-                data_dict["x"],
-                data_dict["y"],
-                data_dict["z"],
-            )
+        # TODO: This is a hack, should be simplified in the future
+        self.model = model
+
+        def helmholtz(data_dict: Dict[str, paddle.Tensor]) -> paddle.Tensor:
+            xs = tuple(data_dict[invar.name] for invar in invars)
 
             # TODO: Hard code here, for hvp_revrev requires tuple input(s) but not dict
-            u__x__x = hvp_revrev(lambda x_: self.model.forward_tensor(x_, y, z), (x,))
-            u__y__y = hvp_revrev(lambda y_: self.model.forward_tensor(x, y_, z), (y,))
-            u__z__z = hvp_revrev(lambda z_: self.model.forward_tensor(x, y, z_), (z,))
+            if self.dim == 1:
+                u__x__x = hvp_revrev(lambda x_: self.model.forward_tensor(x_), (xs[0],))
+                out = (self.k**2) * data_dict["u"] + u__x__x
+            elif self.dim == 2:
+                u__x__x = hvp_revrev(
+                    lambda x_: self.model.forward_tensor(x_, xs[1]), (xs[0],)
+                )
+                u__y__y = hvp_revrev(
+                    lambda y_: self.model.forward_tensor(xs[0], y_), (xs[1],)
+                )
+                out = (self.k**2) * data_dict["u"] + u__x__x + u__y__y
+            elif self.dim >= 3:
+                u__x__x = hvp_revrev(
+                    lambda x_: self.model.forward_tensor(x_, xs[1], xs[2]), (xs[0],)
+                )
+                u__y__y = hvp_revrev(
+                    lambda y_: self.model.forward_tensor(xs[0], y_, xs[2]), (xs[1],)
+                )
+                u__z__z = hvp_revrev(
+                    lambda z_: self.model.forward_tensor(xs[0], xs[1], z_), (xs[2],)
+                )
+                out = (self.k**2) * data_dict["u"] + u__x__x + u__y__y + u__z__z
+            else:
+                raise NotImplementedError(
+                    f"dim should be less or equal to 3, but got {self.dim}."
+                )
 
-            out = (self.k**2) * data_dict["u"] + u__x__x + u__y__y + u__z__z
             return out
 
         self.add_equation("helmholtz", helmholtz)

test/equation/test_helmholtz.py

@@ -0,0 +1,41 @@
+import numpy as np
+import paddle
+import pytest
+
+from ppsci import arch
+from ppsci import equation
+
+__all__ = []
+
+
+@pytest.mark.parametrize("dim", (1, 2, 3))
+def test_helmholtz(dim):
+    """Test for only mean."""
+    nrs = [3, 4, 5][:dim]
+    input_keys = ("x", "y", "z")[:dim]
+    output_keys = ("u",)
+
+    # generate input data
+    input_dict = {
+        input_keys[i]: paddle.to_tensor(
+            np.random.randn(nrs[i], 1).astype(np.float32), stop_gradient=False
+        )
+        for i in range(dim)
+    }
+    model = arch.SPINN(input_keys, output_keys, r=16, num_layers=2, hidden_size=8)
+    y = model(input_dict)["u"]
+    assert y.shape == [*nrs, 1]
+    data_dict = {
+        **input_dict,
+        "u": y,
+    }
+    helmholtz_obj = equation.Helmholtz(dim, 1.0, model)
+
+    helmholtz_out = helmholtz_obj.equations["helmholtz"](data_dict)
+
+    # check result whether is equal
+    assert helmholtz_out.shape == [*nrs, 1]
+
+
+if __name__ == "__main__":
+    pytest.main()