update features

Author: lijialin03

ppsci/arch/base.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Callable
 from typing import Dict
 from typing import Tuple
 
@@ -80,10 +81,28 @@ def split_to_dict(
         data = paddle.split(data_tensor, len(keys), axis=axis)
         return {key: data[i] for i, key in enumerate(keys)}
 
-    def register_input_transform(self, transform):
+    def register_input_transform(
+        self,
+        transform: Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]],
+    ):
+        """Register input transform.
+
+        Args:
+            transform (Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]]):
+                Input transform of network, receive a single tensor dict and return a single tensor dict.
+        """
         self._input_transform = transform
 
-    def register_output_transform(self, transform):
+    def register_output_transform(
+        self,
+        transform: Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]],
+    ):
+        """Register output transform.
+
+        Args:
+            transform (Callable[[Dict[str, paddle.Tensor]], Dict[str, paddle.Tensor]]):
+                Output transform of network, receive a single tensor dict and return a single tensor dict.
+        """
         self._output_transform = transform
 
     def __str__(self):

ppsci/arch/mlp.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
 from typing import Tuple
 from typing import Union
 
@@ -33,6 +34,7 @@ class MLP(base.Arch):
         activation (str, optional): Name of activation function. Defaults to "tanh".
         skip_connection (bool, optional): Whether to use skip connection. Defaults to False.
         weight_norm (bool, optional): Whether to apply weight norm on parameter(s). Defaults to False.
+        input_dim (Optional[int], optional): Number of input's dimension. Defaults to None.
 
     Examples:
         >>> import ppsci
@@ -48,6 +50,7 @@ def __init__(
         activation: str = "tanh",
         skip_connection: bool = False,
         weight_norm: bool = False,
+        input_dim: Optional[int] = None,
     ):
         super().__init__()
         self.input_keys = input_keys
@@ -71,7 +74,7 @@ def __init__(
             )
 
         # initialize FC layer(s)
-        cur_size = len(self.input_keys)
+        cur_size = len(self.input_keys) if input_dim is None else input_dim
         for _size in hidden_size:
             self.linears.append(nn.Linear(cur_size, _size))
             if weight_norm:

ppsci/constraint/supervised_constraint.py

@@ -72,3 +72,15 @@ def __init__(
 
         # construct dataloader with dataset and dataloader_cfg
         super().__init__(_dataset, dataloader_cfg, loss, name)
+
+    def __str__(self):
+        return ", ".join(
+            [
+                self.__class__.__name__,
+                f"name = {self.name}",
+                f"input_keys = {self.input_keys}",
+                f"output_keys = {self.output_keys}",
+                f"output_expr = {self.output_expr}",
+                f"loss = {self.loss}",
+            ]
+        )

ppsci/data/dataset/mat_dataset.py

@@ -34,7 +34,7 @@ class MatDataset(io.Dataset):
     Args:
         file_path (str): Mat file path.
         input_keys (Tuple[str, ...]): List of input keys.
-        label_keys (Tuple[str, ...]): List of label keys.
+        label_keys (Tuple[str, ...], optional): List of label keys. Defaults to ().
         alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys.
             i.e. {inner_key: outer_key}. Defaults to None.
         weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of
@@ -57,7 +57,7 @@ def __init__(
         self,
         file_path: str,
         input_keys: Tuple[str, ...],
-        label_keys: Tuple[str, ...],
+        label_keys: Tuple[str, ...] = (),
         alias_dict: Optional[Dict[str, str]] = None,
         weight_dict: Optional[Dict[str, Union[Callable, float]]] = None,
         timestamps: Optional[Tuple[float, ...]] = None,
@@ -153,7 +153,7 @@ class IterableMatDataset(io.IterableDataset):
     Args:
         file_path (str): Mat file path.
         input_keys (Tuple[str, ...]): List of input keys.
-        label_keys (Tuple[str, ...]): List of label keys.
+        label_keys (Tuple[str, ...], optional): List of label keys. Defaults to ().
         alias_dict (Optional[Dict[str, str]]): Dict of alias(es) for input and label keys.
             i.e. {inner_key: outer_key}. Defaults to None.
         weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the weight of
@@ -176,7 +176,7 @@ def __init__(
         self,
         file_path: str,
         input_keys: Tuple[str, ...],
-        label_keys: Tuple[str, ...],
+        label_keys: Tuple[str, ...] = (),
         alias_dict: Optional[Dict[str, str]] = None,
         weight_dict: Optional[Dict[str, Union[Callable, float]]] = None,
         timestamps: Optional[Tuple[Union[int, float], ...]] = None,

ppsci/loss/integral.py

@@ -54,7 +54,7 @@ def forward(self, output_dict, label_dict, weight_dict=None):
                 label_dict[key],
                 "none",
             )
-            if weight_dict is not None:
+            if weight_dict:
                 loss *= weight_dict[key]
 
             if self.reduction == "sum":

ppsci/loss/l1.py

@@ -48,7 +48,7 @@ def forward(self, output_dict, label_dict, weight_dict=None):
         losses = 0.0
         for key in label_dict:
             loss = F.l1_loss(output_dict[key], label_dict[key], "none")
-            if weight_dict is not None:
+            if weight_dict:
                 loss *= weight_dict[key]
             if "area" in output_dict:
                 loss *= output_dict["area"]

ppsci/loss/l2.py

@@ -73,7 +73,7 @@ def forward(self, output_dict, label_dict, weight_dict=None):
             loss = F.mse_loss(
                 output_dict[key][:n_output], output_dict[key][n_output:], "none"
             )
-            if weight_dict is not None:
+            if weight_dict:
                 loss *= weight_dict[key]
             if "area" in output_dict:
                 loss *= output_dict["area"]

ppsci/loss/mse.py

@@ -53,7 +53,7 @@ def forward(self, output_dict, label_dict, weight_dict=None):
         losses = 0.0
         for key in label_dict:
             loss = F.mse_loss(output_dict[key], label_dict[key], "none")
-            if weight_dict is not None:
+            if weight_dict:
                 loss *= weight_dict[key]
             if "area" in output_dict:
                 loss *= output_dict["area"]

ppsci/validate/sup_validator.py

@@ -83,3 +83,17 @@ def __init__(
 
         # construct dataloader with dataset and dataloader_cfg
         super().__init__(_dataset, dataloader_cfg, loss, metric, name)
+
+    def __str__(self):
+        return ", ".join(
+            [
+                self.__class__.__name__,
+                f"name = {self.name}",
+                f"input_keys = {self.input_keys}",
+                f"output_keys = {self.output_keys}",
+                f"output_expr = {self.output_expr}",
+                f"len(dataloader) = {len(self.data_loader)}",
+                f"loss = {self.loss}",
+                f"metric = {list(self.metric.keys())}",
+            ]
+        )