[Feature] Replace python funcion with sympy expression (#507)

* add sympy2func

* fix node

* fix var name

* fix implementation

* add sympy to func module

* add type hint and clean code

* update unitest for N-S equation with sympy-base 'nu'

* delete redundant sympy2func.py

* remove sympy to function code from expression.py to sym_to_func.py

* update type hint in expression

* update code

* update code

* update code

* refine sym_to_func.py

* replace sympy PDE for biharmonic and laplace

* refine sym_to_func.py

* fix bug in _cvt_to_key

* refine sym_to_func and expression code

* add euler_beam static code(WIP, can not running, to be debug)

* wip code

* temporary code(need to be refined)

* update solver code

* replace more pdes with sympy

* simplify code in solver

* update code

* rename 'normal_dot_vel' to 'normal_dot_vec'

* fix bug

* update unitest

* remove redundant unitest

* remove unnecessary code

* remove unnecessary more code

* use DETACH_FUNC_NAME instead of 'detach'

* add derivatives for sdf function

* replace .diff.diff with .diff(, 2)

* support exporting expression to .dot and .png file for visualizing and DEBUG

* remove compute_sdf_derivatives for next PR

* refine docstring of ppsci/data/dataset/array_dataset.py

* remove sdf_derivatives code in geometry for next PR

* remove print code in solver

* rename sympy_to_function to lambdify and add it in ppsci.*

* rename for test files

* rename sym_to_func.py to symbolic.py

* update linear_init_ and conv_init_ to kaiming style

* refine probability document

* change list to tuple

* update docstrings of equations

* larger atol to 1e-7 for test_linear_elasticity

* fix seed to 42 for test_linear_elasticity

---------

Co-authored-by: PuQing <[email protected]>

Author: HydrogenSulfate

docs/zh/api/probability.md

@@ -1,3 +1,5 @@
+# Probability(概率编程) 模块
+
 ::: ppsci.probability
     handler: python
     options:

docs/zh/api/utils.md

@@ -19,5 +19,6 @@
         - load_checkpoint
         - load_pretrain
         - save_checkpoint
+        - lambdify
       show_root_heading: false
       heading_level: 3

examples/aneurysm/aneurysm.py

@@ -132,7 +132,7 @@ def inlet_w_ref_func(_in):
     )
     igc_outlet = ppsci.constraint.IntegralConstraint(
         equation["NormalDotVec"].equations,
-        {"normal_dot_vel": 2.54},
+        {"normal_dot_vec": 2.54},
         geom["outlet_geo"],
         {
             **train_dataloader_cfg,
@@ -141,12 +141,12 @@ def inlet_w_ref_func(_in):
             "integral_batch_size": 310,
         },
         ppsci.loss.IntegralLoss("sum"),
-        weight_dict={"normal_dot_vel": 0.1},
+        weight_dict={"normal_dot_vec": 0.1},
         name="igc_outlet",
     )
     igc_integral = ppsci.constraint.IntegralConstraint(
         equation["NormalDotVec"].equations,
-        {"normal_dot_vel": -2.54},
+        {"normal_dot_vec": -2.54},
         geom["integral_geo"],
         {
             **train_dataloader_cfg,
@@ -155,7 +155,7 @@ def inlet_w_ref_func(_in):
             "integral_batch_size": 310,
         },
         ppsci.loss.IntegralLoss("sum"),
-        weight_dict={"normal_dot_vel": 0.1},
+        weight_dict={"normal_dot_vec": 0.1},
         name="igc_integral",
     )
     # wrap constraints together

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/laplace/laplace2d.py

@@ -29,7 +29,7 @@
     EVAL_FREQ = 200
 
     # set output directory
-    OUTPUT_DIR = "./output/laplace2d" if not args.output_dir else args.output_dir
+    OUTPUT_DIR = "./output_laplace2d" if not args.output_dir else args.output_dir
     logger.init_logger("ppsci", f"{OUTPUT_DIR}/train.log", "info")
 
     # set model

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

mkdocs.yml

@@ -87,6 +87,7 @@ nav:
       - ppsci.validate: zh/api/validate.md
       - ppsci.visualize: zh/api/visualize.md
       - ppsci.experimental: zh/api/experimental.md
+      - ppsci.probability: zh/api/probability.md
   - 使用指南: zh/user_guide.md
   - 开发与复现指南:
     - 开发指南: zh/development.md

ppsci/__init__.py

@@ -29,6 +29,7 @@
 
 from ppsci.utils.checker import run_check  # isort:skip
 from ppsci.utils.checker import run_check_mesh  # isort:skip
+from ppsci.utils import lambdify  # isort:skip
 
 __all__ = [
     "arch",
@@ -47,4 +48,5 @@
     "experimental",
     "run_check",
     "run_check_mesh",
+    "lambdify",
 ]

ppsci/constraint/boundary_constraint.py

@@ -23,7 +23,6 @@
 
 import numpy as np
 import sympy
-from sympy.parsing import sympy_parser as sp_parser
 from typing_extensions import Literal
 
 from ppsci import geometry
@@ -86,14 +85,12 @@ def __init__(
         weight_dict: Optional[Dict[str, Union[float, Callable]]] = None,
         name: str = "BC",
     ):
-        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)
-
         self.label_dict = label_dict
         self.input_keys = geom.dim_keys
-        self.output_keys = list(label_dict.keys())
+        self.output_keys = tuple(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"]
@@ -137,9 +134,6 @@ def __init__(
         weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
         if weight_dict is not None:
             for key, value in weight_dict.items():
-                if isinstance(value, str):
-                    value = sp_parser.parse_expr(value)
-
                 if isinstance(value, (int, float)):
                     weight[key] = np.full_like(next(iter(label.values())), value)
                 elif isinstance(value, sympy.Basic):

ppsci/constraint/initial_constraint.py

@@ -23,7 +23,6 @@
 
 import numpy as np
 import sympy
-from sympy.parsing import sympy_parser as sp_parser
 from typing_extensions import Literal
 
 from ppsci import geometry
@@ -89,14 +88,12 @@ def __init__(
         weight_dict: Optional[Dict[str, Callable]] = None,
         name: str = "IC",
     ):
-        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)
-
         self.label_dict = label_dict
         self.input_keys = geom.dim_keys
-        self.output_keys = list(label_dict.keys())
+        self.output_keys = tuple(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"]
@@ -117,8 +114,6 @@ def __init__(
         # prepare label
         label = {}
         for key, value in label_dict.items():
-            if isinstance(value, str):
-                value = sp_parser.parse_expr(value)
             if isinstance(value, (int, float)):
                 label[key] = np.full_like(next(iter(input.values())), value)
             elif isinstance(value, sympy.Basic):
@@ -142,8 +137,6 @@ def __init__(
         weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
         if weight_dict is not None:
             for key, value in weight_dict.items():
-                if isinstance(value, str):
-                    value = sp_parser.parse_expr(value)
                 if isinstance(value, (int, float)):
                     weight[key] = np.full_like(next(iter(label.values())), value)
                 elif isinstance(value, sympy.Basic):

ppsci/constraint/integral_constraint.py

@@ -24,7 +24,6 @@
 import numpy as np
 import paddle
 import sympy
-from sympy.parsing import sympy_parser as sp_parser
 from typing_extensions import Literal
 
 from ppsci import geometry
@@ -86,14 +85,12 @@ def __init__(
         weight_dict: Optional[Dict[str, Callable]] = None,
         name: str = "IgC",
     ):
-        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)
-
         self.label_dict = label_dict
         self.input_keys = geom.dim_keys
-        self.output_keys = list(label_dict.keys())
+        self.output_keys = tuple(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"]
@@ -149,9 +146,6 @@ def __init__(
         weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
         if weight_dict is not None:
             for key, value in weight_dict.items():
-                if isinstance(value, str):
-                    value = sp_parser.parse_expr(value)
-
                 if isinstance(value, (int, float)):
                     weight[key] = np.full_like(next(iter(label.values())), value)
                 elif isinstance(value, sympy.Basic):

ppsci/constraint/interior_constraint.py

@@ -23,7 +23,6 @@
 
 import numpy as np
 import sympy
-from sympy.parsing import sympy_parser as sp_parser
 from typing_extensions import Literal
 
 from ppsci import geometry
@@ -86,14 +85,12 @@ def __init__(
         weight_dict: Optional[Dict[str, Union[Callable, float]]] = None,
         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)
-
         self.label_dict = label_dict
         self.input_keys = geom.dim_keys
-        self.output_keys = list(label_dict.keys())
+        self.output_keys = tuple(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"]
@@ -114,8 +111,6 @@ def __init__(
         # prepare label
         label = {}
         for key, value in label_dict.items():
-            if isinstance(value, str):
-                value = sp_parser.parse_expr(value)
             if isinstance(value, (int, float)):
                 label[key] = np.full_like(next(iter(input.values())), value)
             elif isinstance(value, sympy.Basic):

ppsci/constraint/periodic_constraint.py

@@ -24,7 +24,6 @@
 import numpy as np
 import paddle
 import sympy
-from sympy.parsing import sympy_parser as sp_parser
 from typing_extensions import Literal
 
 from ppsci import geometry
@@ -73,13 +72,11 @@ def __init__(
         weight_dict: Optional[Dict[str, Callable]] = None,
         name: str = "PeriodicBC",
     ):
-        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)
-
         self.input_keys = geom.dim_keys
-        self.output_keys = list(output_expr.keys())
+        self.output_keys = tuple(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"]
@@ -143,9 +140,6 @@ def __init__(
         weight = {key: np.ones_like(next(iter(label.values()))) for key in label}
         if weight_dict is not None:
             for key, value in weight_dict.items():
-                if isinstance(value, str):
-                    value = sp_parser.parse_expr(value)
-
                 if isinstance(value, (int, float)):
                     weight[key] = np.full_like(next(iter(label.values())), value)
                 elif isinstance(value, sympy.Basic):

ppsci/constraint/supervised_constraint.py

@@ -60,19 +60,20 @@ def __init__(
         output_expr: Optional[Dict[str, Callable]] = None,
         name: str = "Sup",
     ):
-        self.output_expr = output_expr
-
         # build dataset
         _dataset = dataset.build_dataset(dataloader_cfg["dataset"])
 
         self.input_keys = _dataset.input_keys
         self.output_keys = (
-            list(output_expr.keys()) if output_expr is not None else _dataset.label_keys
+            tuple(output_expr.keys())
+            if output_expr is not None
+            else _dataset.label_keys
         )
 
+        self.output_expr = output_expr
         if self.output_expr is None:
             self.output_expr = {
-                key: lambda out, k=key: out[k] for key in self.output_keys
+                key: (lambda out, k=key: out[k]) for key in self.output_keys
             }
 
         # construct dataloader with dataset and dataloader_cfg

ppsci/data/dataset/array_dataset.py

@@ -29,9 +29,9 @@ class NamedArrayDataset(io.Dataset):
     Args:
         input (Dict[str, np.ndarray]): Input dict.
         label (Dict[str, np.ndarray]): Label dict.
-        weight (Optional[Dict[str, np.ndarray]], optional): Weight dict.
-        transforms (Optional[vision.Compose], optional): Compose object contains sample wise
-            transform(s).
+        weight (Optional[Dict[str, np.ndarray]]): Weight dict. Defaults to None.
+        transforms (Optional[vision.Compose]): Compose object contains sample wise
+            transform(s). Defaults to None.
 
     Examples:
         >>> import ppsci