support default config content in config module and remove deprecated AttrDict series code

HydrogenSulfate · HydrogenSulfate · commit 2fea541be4de · 2024-05-15T06:23:32.000Z
diff --git a/ppsci/arch/__init__.py b/ppsci/arch/__init__.py
@@ -81,7 +81,7 @@ def build_model(cfg):
     """Build model
 
     Args:
-        cfg (AttrDict): Arch config.
+        cfg (DictConfig): Arch config.
 
     Returns:
         nn.Layer: Model.
diff --git a/ppsci/arch/phycrnet.py b/ppsci/arch/phycrnet.py
@@ -147,7 +147,7 @@ def __init__(
         )
 
         # ConvLSTM
-        self.convlstm = paddle.nn.LayerList(
+        self.ConvLSTM = paddle.nn.LayerList(
             [
                 ConvLSTMCell(
                     input_channels=self.input_channels[i],
@@ -194,16 +194,16 @@ def forward(self, x):
                 x = encoder(x)
 
             # convlstm
-            for i, lstm in enumerate(self.convlstm, self.num_encoder):
+            for i, LSTM in enumerate(self.ConvLSTM):
                 if step == 0:
-                    (h, c) = lstm.init_hidden_tensor(
+                    (h, c) = LSTM.init_hidden_tensor(
                         prev_state=self.initial_state[i - self.num_encoder]
                     )
                     internal_state.append((h, c))
 
                 # one-step forward
                 (h, c) = internal_state[i - self.num_encoder]
-                x, new_c = lstm(x, h, c)
+                x, new_c = LSTM(x, h, c)
                 internal_state[i - self.num_encoder] = (x, new_c)
 
             # output
diff --git a/ppsci/constraint/__init__.py b/ppsci/constraint/__init__.py
@@ -42,7 +42,7 @@ def build_constraint(cfg, equation_dict, geom_dict):
     """Build constraint(s).
 
     Args:
-        cfg (List[AttrDict]): Constraint config list.
+        cfg (List[DictConfig]): Constraint config list.
         equation_dict (Dct[str, Equation]): Equation(s) in dict.
         geom_dict (Dct[str, Geometry]): Geometry(ies) in dict.
 
diff --git a/ppsci/data/dataset/__init__.py b/ppsci/data/dataset/__init__.py
@@ -78,7 +78,7 @@ def build_dataset(cfg) -> "io.Dataset":
     """Build dataset
 
     Args:
-        cfg (List[AttrDict]): dataset config list.
+        cfg (List[DictConfig]): dataset config list.
 
     Returns:
         Dict[str, io.Dataset]: dataset.
diff --git a/ppsci/equation/__init__.py b/ppsci/equation/__init__.py
@@ -54,7 +54,7 @@ def build_equation(cfg):
     """Build equation(s)
 
     Args:
-        cfg (List[AttrDict]): Equation(s) config list.
+        cfg (List[DictConfig]): Equation(s) config list.
 
     Returns:
         Dict[str, Equation]: Equation(s) in dict.
diff --git a/ppsci/equation/pde/base.py b/ppsci/equation/pde/base.py
@@ -22,7 +22,7 @@
 from typing import Union
 
 import paddle
-import sympy as sp
+import sympy
 from paddle import nn
 
 DETACH_FUNC_NAME = "detach"
@@ -33,7 +33,7 @@ class PDE:
 
     def __init__(self):
         super().__init__()
-        self.equations: Dict[str, Union[Callable, sp.Basic]] = {}
+        self.equations = {}
         # for PDE which has learnable parameter(s)
         self.learnable_parameters = nn.ParameterList()
 
@@ -42,7 +42,7 @@ def __init__(self):
     @staticmethod
     def create_symbols(
         symbol_str: str,
-    ) -> Union[sp.Symbol, Tuple[sp.Symbol, ...]]:
+    ) -> Union[sympy.Symbol, Tuple[sympy.Symbol, ...]]:
         """create symbolic variables.
 
         Args:
@@ -61,9 +61,11 @@ def create_symbols(
             >>> print(symbols_xyz)
             (x, y, z)
         """
-        return sp.symbols(symbol_str)
+        return sympy.symbols(symbol_str)
 
-    def create_function(self, name: str, invars: Tuple[sp.Symbol, ...]) -> sp.Function:
+    def create_function(
+        self, name: str, invars: Tuple[sympy.Symbol, ...]
+    ) -> sympy.Function:
         """Create named function depending on given invars.
 
         Args:
@@ -84,73 +86,14 @@ def create_function(self, name: str, invars: Tuple[sp.Symbol, ...]) -> sp.Functi
             >>> print(f)
             f(x, y, z)
         """
-        expr = sp.Function(name)(*invars)
+        expr = sympy.Function(name)(*invars)
 
+        # wrap `expression(...)` to `detach(expression(...))`
+        # if name of expression is in given detach_keys
+        if self.detach_keys and name in self.detach_keys:
+            expr = sympy.Function(DETACH_FUNC_NAME)(expr)
         return expr
 
-    def _apply_detach(self):
-        """
-        Wrap detached sub_expr into detach(sub_expr) to prevent gradient
-        back-propagation, only for those items speicified in self.detach_keys.
-
-        NOTE: This function is expected to be called after self.equations is ready in PDE.__init__.
-
-        Examples:
-            >>> import ppsci
-            >>> ns = ppsci.equation.NavierStokes(1.0, 1.0, 2, False)
-            >>> print(ns)
-            NavierStokes
-                continuity: Derivative(u(x, y), x) + Derivative(v(x, y), y)
-                momentum_x: u(x, y)*Derivative(u(x, y), x) + v(x, y)*Derivative(u(x, y), y) + 1.0*Derivative(p(x, y), x) - 1.0*Derivative(u(x, y), (x, 2)) - 1.0*Derivative(u(x, y), (y, 2))
-                momentum_y: u(x, y)*Derivative(v(x, y), x) + v(x, y)*Derivative(v(x, y), y) + 1.0*Derivative(p(x, y), y) - 1.0*Derivative(v(x, y), (x, 2)) - 1.0*Derivative(v(x, y), (y, 2))
-            >>> detach_keys = ("u", "v__y")
-            >>> ns = ppsci.equation.NavierStokes(1.0, 1.0, 2, False, detach_keys=detach_keys)
-            >>> print(ns)
-            NavierStokes
-                continuity: detach(Derivative(v(x, y), y)) + Derivative(u(x, y), x)
-                momentum_x: detach(u(x, y))*Derivative(u(x, y), x) + v(x, y)*Derivative(u(x, y), y) + 1.0*Derivative(p(x, y), x) - 1.0*Derivative(u(x, y), (x, 2)) - 1.0*Derivative(u(x, y), (y, 2))
-                momentum_y: detach(u(x, y))*Derivative(v(x, y), x) + detach(Derivative(v(x, y), y))*v(x, y) + 1.0*Derivative(p(x, y), y) - 1.0*Derivative(v(x, y), (x, 2)) - 1.0*Derivative(v(x, y), (y, 2))
-        """
-        if self.detach_keys is None:
-            return
-
-        from copy import deepcopy
-
-        from sympy.core.traversal import postorder_traversal
-
-        from ppsci.utils.symbolic import _cvt_to_key
-
-        for name, expr in self.equations.items():
-            if not isinstance(expr, sp.Basic):
-                continue
-            # only process sympy expression
-            expr_ = deepcopy(expr)
-            for item in postorder_traversal(expr):
-                if _cvt_to_key(item) in self.detach_keys:
-                    # inplace all related sub_expr into detach(sub_expr)
-                    expr_ = expr_.replace(item, sp.Function(DETACH_FUNC_NAME)(item))
-
-                    # remove all detach wrapper for more-than-once wrapped items to prevent duplicated wrapping
-                    expr_ = expr_.replace(
-                        sp.Function(DETACH_FUNC_NAME)(
-                            sp.Function(DETACH_FUNC_NAME)(item)
-                        ),
-                        sp.Function(DETACH_FUNC_NAME)(item),
-                    )
-
-                    # remove unccessary detach wrapping for the first arg of Derivative
-                    for item_ in list(postorder_traversal(expr_)):
-                        if isinstance(item_, sp.Derivative):
-                            if item_.args[0].name == DETACH_FUNC_NAME:
-                                expr_ = expr_.replace(
-                                    item_,
-                                    sp.Derivative(
-                                        item_.args[0].args[0], *item_.args[1:]
-                                    ),
-                                )
-
-            self.equations[name] = expr_
-
     def add_equation(self, name: str, equation: Callable):
         """Add an equation.
 
@@ -167,8 +110,7 @@ def add_equation(self, name: str, equation: Callable):
             >>> equation = sympy.diff(u, x) + sympy.diff(u, y)
             >>> pde.add_equation('linear_pde', equation)
             >>> print(pde)
-            PDE
-                linear_pde: 2*x + 2*y
+            PDE, linear_pde: 2*x + 2*y
         """
         self.equations.update({name: equation})
 
@@ -239,7 +181,7 @@ def set_state_dict(
         return self.learnable_parameters.set_state_dict(state_dict)
 
     def __str__(self):
-        return "\n".join(
+        return ", ".join(
             [self.__class__.__name__]
-            + [f"    {name}: {eq}" for name, eq in self.equations.items()]
+            + [f"{name}: {eq}" for name, eq in self.equations.items()]
         )
diff --git a/ppsci/equation/pde/biharmonic.py b/ppsci/equation/pde/biharmonic.py
@@ -70,5 +70,3 @@ def __init__(
                 biharmonic += u.diff(invar_i, 2).diff(invar_j, 2)
 
         self.add_equation("biharmonic", biharmonic)
-
-        self._apply_detach()
diff --git a/ppsci/equation/pde/heat_exchanger.py b/ppsci/equation/pde/heat_exchanger.py
@@ -90,5 +90,3 @@ def __init__(
         self.add_equation("heat_boundary", heat_boundary)
         self.add_equation("cold_boundary", cold_boundary)
         self.add_equation("wall", wall)
-
-        self._apply_detach()
diff --git a/ppsci/equation/pde/laplace.py b/ppsci/equation/pde/laplace.py
@@ -51,5 +51,3 @@ def __init__(self, dim: int, detach_keys: Optional[Tuple[str, ...]] = None):
             laplace += u.diff(invar, 2)
 
         self.add_equation("laplace", laplace)
-
-        self._apply_detach()
diff --git a/ppsci/equation/pde/linear_elasticity.py b/ppsci/equation/pde/linear_elasticity.py
@@ -179,5 +179,3 @@ def __init__(
         self.add_equation("traction_y", traction_y)
         if self.dim == 3:
             self.add_equation("traction_z", traction_z)
-
-        self._apply_detach()
diff --git a/ppsci/equation/pde/navier_stokes.py b/ppsci/equation/pde/navier_stokes.py
@@ -147,5 +147,3 @@ def __init__(
         self.add_equation("momentum_y", momentum_y)
         if self.dim == 3:
             self.add_equation("momentum_z", momentum_z)
-
-        self._apply_detach()
diff --git a/ppsci/equation/pde/nls_m_b.py b/ppsci/equation/pde/nls_m_b.py
@@ -97,5 +97,3 @@ def __init__(
         self.add_equation("Maxwell_1", Maxwell_1)
         self.add_equation("Maxwell_2", Maxwell_2)
         self.add_equation("Bloch", Bloch)
-
-        self._apply_detach()
diff --git a/ppsci/equation/pde/normal_dot_vec.py b/ppsci/equation/pde/normal_dot_vec.py
@@ -55,5 +55,3 @@ def __init__(
             normal_dot_vec += normal * vec
 
         self.add_equation("normal_dot_vec", normal_dot_vec)
-
-        self._apply_detach()
diff --git a/ppsci/equation/pde/poisson.py b/ppsci/equation/pde/poisson.py
@@ -49,5 +49,3 @@ def __init__(self, dim: int, detach_keys: Optional[Tuple[str, ...]] = None):
             poisson += p.diff(invar, 2)
 
         self.add_equation("poisson", poisson)
-
-        self._apply_detach()
diff --git a/ppsci/equation/pde/viv.py b/ppsci/equation/pde/viv.py
@@ -60,5 +60,3 @@ def __init__(self, rho: float, k1: float, k2: float):
         k2 = self.create_symbols(self.k2.name)
         f = self.rho * eta.diff(t_f, 2) + sp.exp(k1) * eta.diff(t_f) + sp.exp(k2) * eta
         self.add_equation("f", f)
-
-        self._apply_detach()
diff --git a/ppsci/geometry/__init__.py b/ppsci/geometry/__init__.py
@@ -54,7 +54,7 @@ def build_geometry(cfg):
     """Build geometry(ies)
 
     Args:
-        cfg (List[AttrDict]): Geometry config list.
+        cfg (List[DictConfig]): Geometry config list.
 
     Returns:
         Dict[str, Geometry]: Geometry(ies) in dict.
diff --git a/ppsci/loss/__init__.py b/ppsci/loss/__init__.py
@@ -53,7 +53,7 @@ def build_loss(cfg):
     """Build loss.
 
     Args:
-        cfg (AttrDict): Loss config.
+        cfg (DictConfig): Loss config.
     Returns:
         Loss: Callable loss object.
     """
diff --git a/ppsci/loss/mtl/__init__.py b/ppsci/loss/mtl/__init__.py
@@ -35,7 +35,7 @@ def build_mtl_aggregator(cfg):
     """Build loss aggregator with multi-task learning method.
 
     Args:
-        cfg (AttrDict): Aggregator config.
+        cfg (DictConfig): Aggregator config.
     Returns:
         Loss: Callable loss aggregator object.
     """
diff --git a/ppsci/metric/__init__.py b/ppsci/metric/__init__.py
@@ -43,7 +43,7 @@ def build_metric(cfg):
     """Build metric.
 
     Args:
-        cfg (List[AttrDict]): List of metric config.
+        cfg (List[DictConfig]): List of metric config.
 
     Returns:
         Dict[str, Metric]: Dict of callable metric object.
diff --git a/ppsci/optimizer/__init__.py b/ppsci/optimizer/__init__.py
@@ -39,7 +39,7 @@ def build_lr_scheduler(cfg, epochs, iters_per_epoch):
     """Build learning rate scheduler.
 
     Args:
-        cfg (AttrDict): Learning rate scheduler config.
+        cfg (DictConfig): Learning rate scheduler config.
         epochs (int): Total epochs.
         iters_per_epoch (int): Number of iterations of one epoch.
 
@@ -57,7 +57,7 @@ def build_optimizer(cfg, model_list, epochs, iters_per_epoch):
     """Build optimizer and learning rate scheduler
 
     Args:
-        cfg (AttrDict): Learning rate scheduler config.
+        cfg (DictConfig): Learning rate scheduler config.
         model_list (Tuple[nn.Layer, ...]): Tuple of model(s).
         epochs (int): Total epochs.
         iters_per_epoch (int): Number of iterations of one epoch.
diff --git a/ppsci/solver/solver.py b/ppsci/solver/solver.py
diff --git a/ppsci/utils/__init__.py b/ppsci/utils/__init__.py
diff --git a/ppsci/utils/callbacks.py b/ppsci/utils/callbacks.py
diff --git a/ppsci/utils/config.py b/ppsci/utils/config.py
diff --git a/ppsci/utils/download.py b/ppsci/utils/download.py
diff --git a/ppsci/utils/symbolic.py b/ppsci/utils/symbolic.py
diff --git a/ppsci/validate/__init__.py b/ppsci/validate/__init__.py
diff --git a/ppsci/visualize/__init__.py b/ppsci/visualize/__init__.py
