update code

examples/bracket/bracket.py

@@ -127,15 +127,15 @@
     support_interior_constraint = ppsci.constraint.InteriorConstraint(
         equation["LinearElasticity"].equations,
         {
-            "equilibrium_x": 0,
-            "equilibrium_y": 0,
-            "equilibrium_z": 0,
             "stress_disp_xx": 0,
             "stress_disp_yy": 0,
             "stress_disp_zz": 0,
             "stress_disp_xy": 0,
             "stress_disp_xz": 0,
             "stress_disp_yz": 0,
+            "equilibrium_x": 0,
+            "equilibrium_y": 0,
+            "equilibrium_z": 0,
         },
         geom["geo"],
         {**train_dataloader_cfg, "batch_size": 2048},
@@ -149,30 +149,30 @@
             & (z < BOUNDS_SUPPORT_Z[1])
         ),
         weight_dict={
-            "equilibrium_x": "sdf",
-            "equilibrium_y": "sdf",
-            "equilibrium_z": "sdf",
             "stress_disp_xx": "sdf",
             "stress_disp_yy": "sdf",
             "stress_disp_zz": "sdf",
             "stress_disp_xy": "sdf",
             "stress_disp_xz": "sdf",
             "stress_disp_yz": "sdf",
+            "equilibrium_x": "sdf",
+            "equilibrium_y": "sdf",
+            "equilibrium_z": "sdf",
         },
         name="support_interior",
     )
     bracket_interior_constraint = ppsci.constraint.InteriorConstraint(
         equation["LinearElasticity"].equations,
         {
-            "equilibrium_x": 0,
-            "equilibrium_y": 0,
-            "equilibrium_z": 0,
             "stress_disp_xx": 0,
             "stress_disp_yy": 0,
             "stress_disp_zz": 0,
             "stress_disp_xy": 0,
             "stress_disp_xz": 0,
             "stress_disp_yz": 0,
+            "equilibrium_x": 0,
+            "equilibrium_y": 0,
+            "equilibrium_z": 0,
         },
         geom["geo"],
         {**train_dataloader_cfg, "batch_size": 1024},
@@ -186,15 +186,15 @@
             & (z < BOUNDS_BRACKET_Z[1])
         ),
         weight_dict={
-            "equilibrium_x": "sdf",
-            "equilibrium_y": "sdf",
-            "equilibrium_z": "sdf",
             "stress_disp_xx": "sdf",
             "stress_disp_yy": "sdf",
             "stress_disp_zz": "sdf",
             "stress_disp_xy": "sdf",
             "stress_disp_xz": "sdf",
             "stress_disp_yz": "sdf",
+            "equilibrium_x": "sdf",
+            "equilibrium_y": "sdf",
+            "equilibrium_z": "sdf",
         },
         name="bracket_interior",
     )

examples/pipe/poiseuille_flow.py

@@ -133,9 +133,7 @@ def output_trans_p(input, out):
 
     # set euqation
     equation = {
-        "NavierStokes": ppsci.equation.NavierStokes(
-            nu=lambda out: out["nu"], rho=RHO, dim=2, time=False
-        )
+        "NavierStokes": ppsci.equation.NavierStokes(nu="nu", rho=RHO, dim=2, time=False)
     }
 
     # set constraint

ppsci/constraint/boundary_constraint.py

@@ -87,13 +87,16 @@ def __init__(
         name: str = "BC",
     ):
         self.output_expr = output_expr
-        for label_name, expr in self.output_expr.items():
+        for output_name, expr in self.output_expr.items():
             if isinstance(expr, str):
-                self.output_expr[label_name] = sp_parser.parse_expr(expr)
+                self.output_expr[output_name] = sp_parser.parse_expr(expr)
 
         self.label_dict = label_dict
         self.input_keys = geom.dim_keys
         self.output_keys = list(label_dict.keys())
+        self.output_expr = {
+            k: v for k, v in output_expr.items() if k in self.output_keys
+        }
         # "area" will be kept in "output_dict" for computation.
         if isinstance(geom, geometry.Mesh):
             self.output_keys += ["area"]

ppsci/constraint/initial_constraint.py

@@ -90,13 +90,16 @@ def __init__(
         name: str = "IC",
     ):
         self.output_expr = output_expr
-        for label_name, expr in self.output_expr.items():
+        for output_name, expr in self.output_expr.items():
             if isinstance(expr, str):
-                self.output_expr[label_name] = sp_parser.parse_expr(expr)
+                self.output_expr[output_name] = sp_parser.parse_expr(expr)
 
         self.label_dict = label_dict
         self.input_keys = geom.dim_keys
         self.output_keys = list(label_dict.keys())
+        self.output_expr = {
+            k: v for k, v in output_expr.items() if k in self.output_keys
+        }
         # "area" will be kept in "output_dict" for computation.
         if isinstance(geom.geometry, geometry.Mesh):
             self.output_keys += ["area"]

ppsci/constraint/integral_constraint.py

@@ -87,13 +87,16 @@ def __init__(
         name: str = "IgC",
     ):
         self.output_expr = output_expr
-        for label_name, expr in self.output_expr.items():
+        for output_name, expr in self.output_expr.items():
             if isinstance(expr, str):
-                self.output_expr[label_name] = sp_parser.parse_expr(expr)
+                self.output_expr[output_name] = sp_parser.parse_expr(expr)
 
         self.label_dict = label_dict
         self.input_keys = geom.dim_keys
         self.output_keys = list(label_dict.keys())
+        self.output_expr = {
+            k: v for k, v in output_expr.items() if k in self.output_keys
+        }
         # "area" will be kept in "output_dict" for computation.
         if isinstance(geom, geometry.Mesh):
             self.output_keys += ["area"]

ppsci/constraint/interior_constraint.py

@@ -87,13 +87,16 @@ def __init__(
         name: str = "EQ",
     ):
         self.output_expr = output_expr
-        for label_name, expr in self.output_expr.items():
-            if isinstance(expr, str):
-                self.output_expr[label_name] = sp_parser.parse_expr(expr)
+        for output_name, expr in self.output_expr.items():
+            if isinstance(expr, str) and output_name in label_dict:
+                self.output_expr[output_name] = sp_parser.parse_expr(expr)
 
         self.label_dict = label_dict
         self.input_keys = geom.dim_keys
         self.output_keys = list(label_dict.keys())
+        self.output_expr = {
+            k: v for k, v in output_expr.items() if k in self.output_keys
+        }
         # "area" will be kept in "output_dict" for computation.
         if isinstance(geom, geometry.Mesh):
             self.output_keys += ["area"]

ppsci/constraint/periodic_constraint.py

@@ -74,12 +74,15 @@ def __init__(
         name: str = "PeriodicBC",
     ):
         self.output_expr = output_expr
-        for label_name, expr in self.output_expr.items():
+        for output_name, expr in self.output_expr.items():
             if isinstance(expr, str):
-                self.output_expr[label_name] = sp_parser.parse_expr(expr)
+                self.output_expr[output_name] = sp_parser.parse_expr(expr)
 
         self.input_keys = geom.dim_keys
         self.output_keys = list(output_expr.keys())
+        self.output_expr = {
+            k: v for k, v in output_expr.items() if k in self.output_keys
+        }
         # "area" will be kept in "output_dict" for computation.
         if isinstance(geom, geometry.Mesh):
             self.output_keys += ["area"]

ppsci/equation/pde/base.py

@@ -17,6 +17,7 @@
 from typing import Callable
 from typing import Dict
 from typing import List
+from typing import Optional
 from typing import Tuple
 
 import paddle
@@ -30,12 +31,11 @@ class PDE:
     def __init__(self):
         super().__init__()
         self.equations = {}
-        # self.equations_func = {}
-        # self.detach_keys = []
-
         # for PDE which has learnable parameter(s)
         self.learnable_parameters = nn.ParameterList()
 
+        self.detach_keys: Optional[Tuple[str, ...]] = None
+
     def create_symbols(self, symbol_str: str) -> Tuple[sympy.Symbol, ...]:
         """Create symbols
 
@@ -59,7 +59,24 @@ def create_function(
         Returns:
             sympy.Function: Named sympy function.
         """
-        return sympy.Function(name)(*invars)
+        expr = sympy.Function(name)(*invars)
+        if self.detach_keys and name in self.detach_keys:
+            expr = sympy.Function("detach")(expr)
+        return expr
+
+    def create_detach(
+        self,
+        expr: sympy.Basic,
+    ) -> sympy.Function:
+        """Create detach function for given expression.
+
+        Args:
+            expr (sympy.Basic): Given expression to be detached.
+
+        Returns:
+            sympy.Function: Detached expression.
+        """
+        return sympy.Function("detach")(expr)
 
     def add_equation(self, name: str, equation: Callable):
         """Add an equation.

ppsci/equation/pde/biharmonic.py

@@ -14,6 +14,9 @@
 
 from __future__ import annotations
 
+from typing import Optional
+from typing import Tuple
+
 from ppsci.equation.pde import base
 
 
@@ -34,14 +37,28 @@ class Biharmonic(base.PDE):
         >>> pde = ppsci.equation.Biharmonic(2, -1.0, 1.0)
     """
 
-    def __init__(self, dim: int, q: float, D: float):
+    def __init__(
+        self,
+        dim: int,
+        q: float,
+        D: float,
+        detach_keys: Optional[Tuple[str, ...]] = None,
+    ):
         super().__init__()
+        self.detach_keys = detach_keys
+
+        invars = self.create_symbols("x y z")[:dim]
+        u = self.create_function("u", invars)
+
+        if isinstance(q, str):
+            q = self.create_function("q", invars)
+        if isinstance(D, str):
+            D = self.create_function("D", invars)
+
         self.dim = dim
         self.q = q
         self.D = D
 
-        invars = self.create_symbols("x y z")[: self.dim]
-        u = self.create_function("u", invars)
         biharmonic = -self.q / self.D
         for invar_i in invars:
             for invar_j in invars:

ppsci/equation/pde/laplace.py

@@ -14,6 +14,9 @@
 
 from __future__ import annotations
 
+from typing import Optional
+from typing import Tuple
+
 from ppsci.equation.pde import base
 
 
@@ -32,12 +35,15 @@ class Laplace(base.PDE):
         >>> pde = ppsci.equation.Laplace(2)
     """
 
-    def __init__(self, dim: int):
+    def __init__(self, dim: int, detach_keys: Optional[Tuple[str, ...]] = None):
         super().__init__()
-        self.dim = dim
-        invars = self.create_symbols("x y z")[: self.dim]
+        self.detach_keys = detach_keys
+
+        invars = self.create_symbols("x y z")[:dim]
         u = self.create_function("u", invars)
 
+        self.dim = dim
+
         laplace = 0
         for invar in invars:
             laplace += u.diff(invar).diff(invar)

ppsci/equation/pde/linear_elasticity.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 from typing import Optional
+from typing import Tuple
 
 import sympy as sp
 
@@ -61,8 +62,10 @@ def __init__(
         rho: float = 1,
         dim: int = 3,
         time: bool = False,
+        detach_keys: Optional[Tuple[str, ...]] = None,
     ):
         super().__init__()
+        self.detach_keys = detach_keys
         self.dim = dim
         self.time = time
 

ppsci/equation/pde/navier_stokes.py

@@ -15,10 +15,10 @@
 from __future__ import annotations
 
 from typing import Callable
+from typing import Optional
+from typing import Tuple
 from typing import Union
 
-import sympy as sp
-
 from ppsci.equation.pde import base
 
 
@@ -63,19 +63,37 @@ class NavierStokes(base.PDE):
         >>> pde = ppsci.equation.NavierStokes(0.1, 1.0, 3, False)
     """
 
-    def __init__(self, nu: Union[float, Callable], rho: float, dim: int, time: bool):
+    def __init__(
+        self,
+        nu: Union[float, Callable],
+        rho: float,
+        dim: int,
+        time: bool,
+        detach_keys: Optional[Tuple[str, ...]] = None,
+    ):
         super().__init__()
+        self.detach_keys = detach_keys
+        t, x, y, z = self.create_symbols("t x y z")
+        invars = (x, y)
+        if time:
+            invars = (t,) + invars
+        if dim == 3:
+            invars += (z,)
+
         self.nu = nu
         self.rho = rho
         self.dim = dim
         self.time = time
 
-        t, x, y, z = self.create_symbols("t x y z")
-        u, v, w, p = self.create_symbols("u v w p")
-        if self.dim == 2:
-            w = sp.Number(0)
-        if not time:
-            t = sp.Number(0)
+        if isinstance(nu, str):
+            nu = self.create_function(nu, invars)
+        if isinstance(rho, str):
+            rho = self.create_function(rho, invars)
+
+        u = self.create_function("u", invars)
+        v = self.create_function("v", invars)
+        w = self.create_function("w", invars)
+        p = self.create_function("p", invars)
 
         continuity = u.diff(x) + v.diff(y) + w.diff(z)
         momentum_x = (
@@ -105,4 +123,5 @@ def __init__(self, nu: Union[float, Callable], rho: float, dim: int, time: bool)
         self.add_equation("continuity", continuity)
         self.add_equation("momentum_x", momentum_x)
         self.add_equation("momentum_y", momentum_y)
-        self.add_equation("momentum_z", momentum_z)
+        if self.dim == 3:
+            self.add_equation("momentum_z", momentum_z)