simplify expression_solver for to_static

HydrogenSulfate · HydrogenSulfate · commit 635e300933e6 · 2023-05-17T05:54:41.000Z
diff --git a/ppsci/solver/eval.py b/ppsci/solver/eval.py
@@ -59,15 +59,12 @@ def eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
 
             for v in input_dict.values():
                 v.stop_gradient = False
-            evaluator = expression.ExpressionSolver(
-                _validator.input_keys, _validator.output_keys, solver.model
-            )
-            for output_name, output_formula in _validator.output_expr.items():
-                evaluator.add_target_expr(output_formula, output_name)
 
             # forward
             with solver.autocast_context_manager(), solver.no_grad_context_manager():
-                output_dict = evaluator(input_dict)
+                output_dict = solver.expr_helper(
+                    _validator.output_expr, input_dict, solver.model
+                )
                 validator_loss = _validator.loss(output_dict, label_dict, weight_dict)
 
             loss_dict[f"loss({_validator.name})"] = float(validator_loss)
diff --git a/ppsci/solver/solver.py b/ppsci/solver/solver.py
@@ -36,6 +36,7 @@
 
 import ppsci
 from ppsci.utils import config
+from ppsci.utils import expression
 from ppsci.utils import logger
 from ppsci.utils import misc
 from ppsci.utils import save_load
@@ -246,6 +247,8 @@ def __init__(
         if logger._logger is not None:
             logger.info(f"Using paddlepaddle {paddle_version} on device {self.device}")
 
+        self.expr_helper = expression.ExpressionSolver()
+
     @staticmethod
     def from_config(cfg: Dict[str, Any]) -> Solver:
         """Initialize solver from given config.
diff --git a/ppsci/solver/train.py b/ppsci/solver/train.py
@@ -52,15 +52,12 @@ def train_epoch_func(solver, epoch_id: int, log_freq: int):
 
             for v in input_dict.values():
                 v.stop_gradient = False
-            evaluator = expression.ExpressionSolver(
-                _constraint.input_keys, _constraint.output_keys, solver.model
-            )
-            for output_name, output_formula in _constraint.output_expr.items():
-                evaluator.add_target_expr(output_formula, output_name)
 
             # forward for every constraint
             with solver.autocast_context_manager():
-                output_dict = evaluator(input_dict)
+                output_dict = solver.expr_helper(
+                    _constraint.output_expr, input_dict, solver.model
+                )
                 constraint_loss = _constraint.loss(output_dict, label_dict, weight_dict)
                 total_loss += constraint_loss
 
diff --git a/ppsci/solver/visu.py b/ppsci/solver/visu.py
@@ -53,15 +53,11 @@ def visualize_func(solver, epoch_id: int):
                     batch_input_dict[key] = input_dict[key][st:ed]
                 batch_input_dict[key].stop_gradient = False
 
-            evaluator = expression.ExpressionSolver(
-                _visualizer.input_keys, _visualizer.output_keys, solver.model
-            )
-            for output_key, output_expr in _visualizer.output_expr.items():
-                evaluator.add_target_expr(output_expr, output_key)
-
             # forward
             with solver.autocast_context_manager():
-                batch_output_dict = evaluator(batch_input_dict)
+                batch_output_dict = solver.expr_helper(
+                    _visualizer.output_expr, batch_input_dict, solver.model
+                )
 
             # collect batch data
             for key, batch_input in batch_input_dict.items():
diff --git a/ppsci/utils/expression.py b/ppsci/utils/expression.py
@@ -15,13 +15,9 @@
 from typing import Callable
 from typing import Union
 
-import paddle
-import sympy
 from paddle import nn
 
 from ppsci.autodiff import clear
-from ppsci.autodiff import hessian
-from ppsci.autodiff import jacobian
 
 
 class ExpressionSolver(nn.Layer):
@@ -38,127 +34,25 @@ class ExpressionSolver(nn.Layer):
         >>> expr_solver = ExpressionSolver(("x", "y"), ("u", "v"), model)
     """
 
-    def __init__(self, input_keys, output_keys, model):
+    def __init__(self):
         super().__init__()
-        self.input_keys = input_keys
-        self.output_keys = output_keys
-        self.model = model
-        self.expr_dict = {}
-        self.output_dict = {}
-
-    def solve_expr(self, expr: sympy.Basic) -> Union[float, paddle.Tensor]:
-        """Evaluates the value of the expression recursively in the expression tree
-            by post-order traversal.
-
-        Args:
-            expr (sympy.Basic): Expression.
-
-        Returns:
-            Union[float, paddle.Tensor]: Value of current expression `expr`.
-        """
-        # already computed in output_dict(including input data)
-        if getattr(expr, "name", None) in self.output_dict:
-            return self.output_dict[expr.name]
-
-        # compute output from model
-        if isinstance(expr, sympy.Symbol):
-            if expr.name in self.model.output_keys:
-                out_dict = self.model(self.output_dict)
-                self.output_dict.update(out_dict)
-                return self.output_dict[expr.name]
-            else:
-                raise ValueError(f"varname {expr.name} not exist!")
-
-        # compute output from model
-        elif isinstance(expr, sympy.Function):
-            out_dict = self.model(self.output_dict)
-            self.output_dict.update(out_dict)
-            return self.output_dict[expr.name]
-
-        # compute derivative
-        elif isinstance(expr, sympy.Derivative):
-            ys = self.solve_expr(expr.args[0])
-            ys_name = expr.args[0].name
-            if ys_name not in self.output_dict:
-                self.output_dict[ys_name] = ys
-            xs = self.solve_expr(expr.args[1][0])
-            xs_name = expr.args[1][0].name
-            if xs_name not in self.output_dict:
-                self.output_dict[xs_name] = xs
-            order = expr.args[1][1]
-            if order == 1:
-                der = jacobian(self.output_dict[ys_name], self.output_dict[xs_name])
-                der_name = f"{ys_name}__{xs_name}"
-            elif order == 2:
-                der = hessian(self.output_dict[ys_name], self.output_dict[xs_name])
-                der_name = f"{ys_name}__{xs_name}__{xs_name}"
-            else:
-                raise NotImplementedError(
-                    f"Expression {expr} has derivative order({order}) >=3, "
-                    f"which is not implemented yet"
-                )
-            if der_name not in self.output_dict:
-                self.output_dict[der_name] = der
-            return der
-
-        # return single python number directly for leaf node
-        elif isinstance(expr, sympy.Number):
-            return float(expr)
-
-        # compute sub-nodes value and merge by addition
-        elif isinstance(expr, sympy.Add):
-            results = [self.solve_expr(arg) for arg in expr.args]
-            out = results[0]
-            for i in range(1, len(results)):
-                out = out + results[i]
-            return out
-
-        # compute sub-nodes value and merge by multiplication
-        elif isinstance(expr, sympy.Mul):
-            results = [self.solve_expr(arg) for arg in expr.args]
-            out = results[0]
-            for i in range(1, len(results)):
-                out = out * results[i]
-            return out
-
-        # compute sub-nodes value and merge by power
-        elif isinstance(expr, sympy.Pow):
-            results = [self.solve_expr(arg) for arg in expr.args]
-            return results[0] ** results[1]
-        else:
-            raise ValueError(
-                f"Expression {expr} of type({type(expr)}) can't be solved yet."
-            )
-
-    def forward(self, input_dict):
-        self.output_dict = input_dict
-        if callable(next(iter(self.expr_dict.values()))):
-            model_output_dict = self.model(input_dict)
-            self.output_dict.update(model_output_dict)
-
-        for name, expr in self.expr_dict.items():
-            if isinstance(expr, sympy.Basic):
-                self.output_dict[name] = self.solve_expr(expr)
-            elif callable(expr):
-                self.output_dict[name] = expr(self.output_dict)
+
+    def forward(self, expr_dict, input_dict, model):
+        output_dict = {k: v for k, v in input_dict.items()}
+
+        # model forward
+        if callable(next(iter(expr_dict.values()))):
+            model_output_dict = model(input_dict)
+            output_dict.update(model_output_dict)
+
+        # equation forward
+        for name, expr in expr_dict.items():
+            if callable(expr):
+                output_dict[name] = expr(output_dict)
             else:
                 raise TypeError(f"expr type({type(expr)}) is invalid")
 
         # clear differentiation cache
         clear()
 
-        return {k: self.output_dict[k] for k in self.output_keys}
-
-    def add_target_expr(self, expr: Callable, expr_name: str):
-        """Add an expression `expr` named `expr_name` to
-
-        Args:
-            expr (Callable): Callable function for computing an expression.
-            expr_name (str): Name of expression.
-        """
-        self.expr_dict[expr_name] = expr
-
-    def __str__(self):
-        return f"input: {self.input_keys}, output: {self.output_keys}\n" + "\n".join(
-            [f"{name} = {expr}" for name, expr in self.expr_dict.items()]
-        )
+        return output_dict
