Merge pull request #328 from HydrogenSulfate/refine_code

HydrogenSulfate · web-flow · commit b7ad945a5ff9 · 2023-05-22T20:35:35.000+08:00
optimize code
diff --git a/docs/zh/api/data/dataset.md b/docs/zh/api/data/dataset.md
@@ -6,11 +6,14 @@
       members:
         - IterableNamedArrayDataset
         - NamedArrayDataset
-        - CylinderDataset
-        - LorenzDataset
-        - RosslerDataset
         - CSVDataset
-        - MatDataset
+        - IterableCSVDataset
         - ERA5Dataset
         - ERA5SampledDataset
+        - IterableMatDataset
+        - MatDataset
+        - CylinderDataset
+        - LorenzDataset
+        - RosslerDataset
+        - VtuDataset
       show_root_heading: false
diff --git a/docs/zh/api/equation.md b/docs/zh/api/equation.md
@@ -10,5 +10,6 @@
         - NavierStokes
         - NormalDotVec
         - Poisson
+        - Vibration
       show_root_heading: false
       heading_level: 3
diff --git a/docs/zh/api/geometry.md b/docs/zh/api/geometry.md
@@ -20,25 +20,3 @@
         - TimeXGeometry
       show_root_heading: false
       heading_level: 3
-
-<!-- # Geometry
-
-::: ppsci.arch
-
-This is on a separate line
-
-$$
-\operatorname{ker} f=\{g\in G:f(g)=e_{H}\}{\mbox{.}}
-$$
-
-The homomorphism $f$ is injective if and only if its kernel is only the
-singleton set $e_G$, because otherwise $\exists a,b\in G$ with $a\neq b$ such
-that $f(a)=f(b)$.
-
-```python
-
---8<--
-./ppsci/data/dataset/array_dataset.py:16:49
---8<--
-
-``` -->
diff --git a/docs/zh/api/lr_scheduler.md b/docs/zh/api/lr_scheduler.md
@@ -4,7 +4,6 @@
     handler: python
     options:
       members:
-        - ConstLR
         - Linear
         - Cosine
         - Step
diff --git a/docs/zh/api/metric.md b/docs/zh/api/metric.md
@@ -4,6 +4,7 @@
     handler: python
     options:
       members:
+        - Metric
         - MAE
         - MSE
         - RMSE
diff --git a/docs/zh/api/utils.md b/docs/zh/api/utils.md
@@ -7,7 +7,9 @@
         - initializer
         - logger
         - misc
-        - reader
+        - load_csv_file
+        - load_mat_file
+        - load_vtk_file
         - run_check
         - AttrDict
         - ExpressionSolver
diff --git a/examples/cylinder/3d_unsteady_discrete/cylinder3d_unsteady.py b/examples/cylinder/3d_unsteady_discrete/cylinder3d_unsteady.py
@@ -342,7 +342,6 @@
             600000,
             label,
             time_list,
-            len(time_list),
             "result_uvwp",
         )
     }
diff --git a/ppsci/arch/base.py b/ppsci/arch/base.py
@@ -62,6 +62,8 @@ def concat_to_tensor(
         Returns:
             Tuple[paddle.Tensor, ...]: Concatenated tensor.
         """
+        if len(keys) == 1:
+            return data_dict[keys[0]]
         data = [data_dict[key] for key in keys]
         return paddle.concat(data, axis)
 
@@ -78,6 +80,8 @@ def split_to_dict(
         Returns:
             Dict[str, paddle.Tensor]: Dict contains tensor.
         """
+        if len(keys) == 1:
+            return {keys[0]: data_tensor}
         data = paddle.split(data_tensor, len(keys), axis=axis)
         return {key: data[i] for i, key in enumerate(keys)}
 
diff --git a/ppsci/data/__init__.py b/ppsci/data/__init__.py
@@ -83,10 +83,9 @@ def build_dataloader(_dataset, cfg):
     # build collate_fn if specified
     batch_transforms_cfg = cfg.pop("batch_transforms", None)
 
+    collate_fn = None
     if isinstance(batch_transforms_cfg, dict) and batch_transforms_cfg:
         collate_fn = batch_transform.build_batch_transforms(batch_transforms_cfg)
-    else:
-        collate_fn = batch_transform.default_collate_fn_allow_none
 
     # build init function
     init_fn = partial(
@@ -97,7 +96,7 @@ def build_dataloader(_dataset, cfg):
     )
 
     # build dataloader
-    dataloader = io.DataLoader(
+    dataloader_ = io.DataLoader(
         dataset=_dataset,
         places=device.get_device(),
         batch_sampler=sampler,
@@ -107,4 +106,4 @@ def build_dataloader(_dataset, cfg):
         worker_init_fn=init_fn,
     )
 
-    return dataloader
+    return dataloader_
diff --git a/ppsci/data/dataset/array_dataset.py b/ppsci/data/dataset/array_dataset.py
@@ -56,11 +56,7 @@ def __init__(
     def __getitem__(self, idx):
         input_item = {key: value[idx] for key, value in self.input.items()}
         label_item = {key: value[idx] for key, value in self.label.items()}
-        weight_item = (
-            {key: value[idx] for key, value in self.weight.items()}
-            if self.weight is not None
-            else None
-        )
+        weight_item = {key: value[idx] for key, value in self.weight.items()}
 
         # TODO(sensen): Transforms may be applied on label and weight.
         if self.transforms is not None:
diff --git a/ppsci/data/dataset/era5_dataset.py b/ppsci/data/dataset/era5_dataset.py
@@ -68,14 +68,12 @@ def __init__(
         self.label_keys = label_keys
         self.precip_file_path = precip_file_path
 
-        self.weight_dict = weight_dict
+        self.weight_dict = {} if weight_dict is None else weight_dict
         if weight_dict is not None:
             self.weight_dict = {key: 1.0 for key in self.label_keys}
             self.weight_dict.update(weight_dict)
 
-        self.vars_channel = (
-            vars_channel if vars_channel is not None else [i for i in range(20)]
-        )
+        self.vars_channel = list(range(20)) if vars_channel is None else vars_channel
         self.num_label_timestamps = num_label_timestamps
         self.transforms = transforms
         self.training = training
@@ -127,6 +125,7 @@ def __getitem__(self, global_idx):
             input_idx, label_idx = local_idx, local_idx + step
 
         input_item = {self.input_keys[0]: input_file[input_idx, self.vars_channel]}
+
         label_item = {}
         for i in range(self.num_label_timestamps):
             if self.precip_file_path is not None:
@@ -138,14 +137,11 @@ def __getitem__(self, global_idx):
                     label_idx + i, self.vars_channel
                 ]
 
-        if self.weight_dict is not None:
-            weight_shape = [1] * len(next(iter(label_item.values)).shape)
-            weight_item = {
-                key: np.full(weight_shape, value, paddle.get_default_dtype())
-                for key, value in self.weight_dict.items()
-            }
-        else:
-            weight_item = None
+        weight_shape = [1] * len(next(iter(label_item.values)).shape)
+        weight_item = {
+            key: np.full(weight_shape, value, paddle.get_default_dtype())
+            for key, value in self.weight_dict.items()
+        }
 
         if self.transforms is not None:
             input_item, label_item, weight_item = self.transforms(
@@ -187,7 +183,7 @@ def __init__(
         self.input_keys = input_keys
         self.label_keys = label_keys
 
-        self.weight_dict = weight_dict
+        self.weight_dict = {} if weight_dict is None else weight_dict
         if weight_dict is not None:
             self.weight_dict = {key: 1.0 for key in self.label_keys}
             self.weight_dict.update(weight_dict)
@@ -201,8 +197,8 @@ def read_data(self, path: str):
         paths = glob.glob(path + "/*.h5")
         paths.sort()
         files = []
-        for path in paths:
-            _file = h5py.File(path, "r")
+        for _path in paths:
+            _file = h5py.File(_path, "r")
             files.append(_file)
         return files
 
@@ -217,20 +213,18 @@ def __getitem__(self, global_idx):
             input_item[key] = np.asarray(
                 _file["input_dict"][key], paddle.get_default_dtype()
             )
+
         label_item = {}
         for key in _file["label_dict"]:
             label_item[key] = np.asarray(
                 _file["label_dict"][key], paddle.get_default_dtype()
             )
 
-        if self.weight_dict is not None:
-            weight_shape = [1] * len(next(iter(label_item.values)).shape)
-            weight_item = {
-                key: np.full(weight_shape, value, paddle.get_default_dtype())
-                for key, value in self.weight_dict.items()
-            }
-        else:
-            weight_item = None
+        weight_shape = [1] * len(next(iter(label_item.values)).shape)
+        weight_item = {
+            key: np.full(weight_shape, value, paddle.get_default_dtype())
+            for key, value in self.weight_dict.items()
+        }
 
         if self.transforms is not None:
             input_item, label_item, weight_item = self.transforms(
diff --git a/ppsci/data/dataset/vtu_dataset.py b/ppsci/data/dataset/vtu_dataset.py
@@ -83,7 +83,7 @@ def __init__(
     def __getitem__(self, idx):
         input_item = {key: value[idx] for key, value in self.input.items()}
         label_item = {key: value[idx] for key, value in self.label.items()}
-        return (input_item, label_item, None)
+        return (input_item, label_item, {})
 
     def __len__(self):
         return self.num_samples
diff --git a/ppsci/data/process/batch_transform/__init__.py b/ppsci/data/process/batch_transform/__init__.py
@@ -25,13 +25,13 @@
 
 from ppsci.data.process import transform
 
-__all__ = ["build_batch_transforms", "default_collate_fn_allow_none"]
+__all__ = ["build_batch_transforms"]
 
 
-def default_collate_fn_allow_none(batch: List[Any]) -> Any:
-    """Modified collate function to allow some fields to be None, such as weight field.
+def default_collate_fn(batch: List[Any]) -> Any:
+    """Default_collate_fn for paddle dataloader.
 
-    ref: https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/fluid/dataloader/collate.py#L24
+    ref: https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/io/dataloader/collate.py#L25
 
     Args:
         batch (List[Any]): Batch of samples to be collated.
@@ -40,11 +40,6 @@ def default_collate_fn_allow_none(batch: List[Any]) -> Any:
         Any: Collated batch data.
     """
     sample = batch[0]
-
-    # allow field to be None
-    if sample is None:
-        return None
-
     if isinstance(sample, np.ndarray):
         batch = np.stack(batch, axis=0)
         return batch
@@ -56,15 +51,12 @@ def default_collate_fn_allow_none(batch: List[Any]) -> Any:
     elif isinstance(sample, (str, bytes)):
         return batch
     elif isinstance(sample, Mapping):
-        return {
-            key: default_collate_fn_allow_none([d[key] for d in batch])
-            for key in sample
-        }
+        return {key: default_collate_fn([d[key] for d in batch]) for key in sample}
     elif isinstance(sample, Sequence):
         sample_fields_num = len(sample)
         if not all(len(sample) == sample_fields_num for sample in iter(batch)):
             raise RuntimeError("fileds number not same among samples in a batch")
-        return [default_collate_fn_allow_none(fields) for fields in zip(*batch)]
+        return [default_collate_fn(fields) for fields in zip(*batch)]
 
     raise TypeError(
         "batch data can only contains: tensor, numpy.ndarray, "
@@ -80,6 +72,6 @@ def collate_fn_batch_transforms(batch: List[Any]):
         # apply batch transform on uncollated data
         batch = batch_transforms(batch)
         # then do collate
-        return default_collate_fn_allow_none(batch)
+        return default_collate_fn(batch)
 
     return collate_fn_batch_transforms
diff --git a/ppsci/data/process/transform/preprocess.py b/ppsci/data/process/transform/preprocess.py
@@ -205,7 +205,7 @@ def __call__(self, data):
                     input_item[key] = value.reshape((B * C, H, W))
                 if value.ndim != 3:
                     raise ValueError(
-                        "Only support squeeze data to ndim=3 now, but got ndim={value.ndim}"
+                        f"Only support squeeze data to ndim=3 now, but got ndim={value.ndim}"
                     )
         if "label" in self.apply_keys:
             for key, value in label_item.items():
@@ -214,6 +214,6 @@ def __call__(self, data):
                     label_item[key] = value.reshape((B * C, H, W))
                 if value.ndim != 3:
                     raise ValueError(
-                        "Only support squeeze data to ndim=3 now, but got ndim={value.ndim}"
+                        f"Only support squeeze data to ndim=3 now, but got ndim={value.ndim}"
                     )
         return input_item, label_item, weight_item
diff --git a/ppsci/solver/eval.py b/ppsci/solver/eval.py
@@ -23,7 +23,7 @@
 from ppsci.utils import profiler
 
 
-def eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
+def _eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
     """Evaluation program by dataset.
 
     Args:
@@ -63,7 +63,8 @@ def eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
                 _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)
+                if output_name in label_dict:
+                    evaluator.add_target_expr(output_formula, output_name)
 
             # forward
             with solver.autocast_context_manager(), solver.no_grad_context_manager():
@@ -147,7 +148,7 @@ def eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
     return target_metric
 
 
-def eval_by_batch(solver, epoch_id: int, log_freq: int) -> float:
+def _eval_by_batch(solver, epoch_id: int, log_freq: int) -> float:
     """Evaluation program by batch.
 
     Args:
@@ -260,5 +261,5 @@ def eval_func(solver, epoch_id: int, log_freq: int) -> float:
         float: Target metric computed during evaluation.
     """
     if solver.compute_metric_by_batch:
-        return eval_by_batch(solver, epoch_id, log_freq)
-    return eval_by_dataset(solver, epoch_id, log_freq)
+        return _eval_by_batch(solver, epoch_id, log_freq)
+    return _eval_by_dataset(solver, epoch_id, log_freq)
diff --git a/ppsci/solver/train.py b/ppsci/solver/train.py
@@ -56,7 +56,8 @@ def train_epoch_func(solver, epoch_id: int, log_freq: int):
                 _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)
+                if output_name in label_dict:
+                    evaluator.add_target_expr(output_formula, output_name)
 
             # forward for every constraint
             with solver.autocast_context_manager():
diff --git a/ppsci/utils/__init__.py b/ppsci/utils/__init__.py
@@ -22,6 +22,10 @@
 from ppsci.utils.expression import ExpressionSolver
 from ppsci.utils.misc import AverageMeter
 from ppsci.utils.misc import set_random_seed
+from ppsci.utils.reader import load_csv_file
+from ppsci.utils.reader import load_mat_file
+from ppsci.utils.reader import load_vtk_file
+from ppsci.utils.reader import load_vtk_with_time_file
 from ppsci.utils.save_load import load_checkpoint
 from ppsci.utils.save_load import load_pretrain
 from ppsci.utils.save_load import save_checkpoint
@@ -31,6 +35,10 @@
     "logger",
     "misc",
     "reader",
+    "load_csv_file",
+    "load_mat_file",
+    "load_vtk_file",
+    "load_vtk_with_time_file",
     "dynamic_import_to_globals",
     "run_check",
     "AttrDict",
diff --git a/ppsci/utils/reader.py b/ppsci/utils/reader.py
diff --git a/ppsci/visualize/plot.py b/ppsci/visualize/plot.py
diff --git a/ppsci/visualize/visualizer.py b/ppsci/visualize/visualizer.py
diff --git a/ppsci/visualize/vtu.py b/ppsci/visualize/vtu.py

Original file line number	Diff line number	Diff line change
`@@ -342,7 +342,6 @@`
`342`	`342`	`600000,`
`343`	`343`	`label,`
`344`	`344`	`time_list,`
`345`		`- len(time_list),`
`346`	`345`	`"result_uvwp",`
`347`	`346`	`)`
`348`	`347`	`}`