update visualize code

Author: HydrogenSulfate

examples/cylinder/3d_unsteady_discrete/cylinder3d_unsteady.py

@@ -342,7 +342,6 @@
             600000,
             label,
             time_list,
-            len(time_list),
             "result_uvwp",
         )
     }

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

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)

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():

ppsci/visualize/plot.py

@@ -20,7 +20,6 @@
 
 import imageio
 import matplotlib
-import matplotlib as mpl
 import numpy as np
 import paddle
 from matplotlib import cm
@@ -174,8 +173,8 @@ def _save_plot_from_2d_array(
     """
 
     plt.close("all")
-    mpl.rcParams["xtick.labelsize"] = 5
-    mpl.rcParams["ytick.labelsize"] = 5
+    matplotlib.rcParams["xtick.labelsize"] = 5
+    matplotlib.rcParams["ytick.labelsize"] = 5
 
     fig, ax = plt.subplots(
         len(visu_keys),
@@ -215,7 +214,7 @@ def _save_plot_from_2d_array(
         ticks = np.linspace(0, 1, 5)
         tickLabels = np.linspace(c_min, c_max, 5)
         tickLabels = [f"{t0:02.2f}" for t0 in tickLabels]
-        cbar = mpl.colorbar.ColorbarBase(
+        cbar = matplotlib.colorbar.ColorbarBase(
             ax_cbar, cmap=plt.get_cmap("inferno"), orientation="vertical", ticks=ticks
         )
         cbar.set_ticklabels(tickLabels, fontsize=5)

ppsci/visualize/visualizer.py

@@ -302,32 +302,21 @@ def __init__(
         batch_size: int = 64,
         label_dict: Optional[Dict[str, np.ndarray]] = None,
         time_list: Optional[Tuple[float, ...]] = None,
-        num_timestamps: int = 1,
         prefix: str = "vtu",
     ):
         self.label = label_dict
         self.time_list = time_list
-        super().__init__(input_dict, output_expr, batch_size, num_timestamps, prefix)
-
-    def save(self, filename: str, data_dict: Dict[str, paddle.Tensor]):
-        """Save points result
-
-        Args:
-            filename (str): Output file name with directory.
-            data_dict (Dict[str, paddle.Tensor]): Predicted result.
-        """
+        super().__init__(input_dict, output_expr, batch_size, len(time_list), prefix)
 
+    def save(self, filename: str, data_dict: Dict[str, np.ndarray]):
         n = int((next(iter(data_dict.values()))).shape[0] / self.num_timestamps)
         coord_keys = [x for x in self.input_dict if x != "t"]
         for i in range(len(self.time_list)):
             vtu.save_vtu_to_mesh(
-                filename=osp.join(filename, f"predict_{i+1}.vtu"),
-                data_dict={
-                    key: (data_dict[key].numpy()[i * n : (i + 1) * n])
-                    for key in data_dict
-                },
-                value_keys=self.output_expr,
-                coord_keys=coord_keys,
+                osp.join(filename, f"predict_{i+1}.vtu"),
+                {key: (data_dict[key][i * n : (i + 1) * n]) for key in data_dict},
+                coord_keys,
+                self.output_keys,
             )
 
 

ppsci/visualize/vtu.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 from typing import Dict
 from typing import Tuple
 
@@ -93,12 +92,18 @@ def _save_vtu_from_array(filename, coord, value, value_keys, num_timestamps=1):
         logger.info(f"Visualization result is saved to {filename}.vtu")
 
 
-def save_vtu_from_dict(filename, data_dict, coord_keys, value_keys, num_timestamps=1):
+def save_vtu_from_dict(
+    filename: str,
+    data_dict: Dict[str, np.ndarray],
+    coord_keys: Tuple[str, ...],
+    value_keys: Tuple[str, ...],
+    num_timestamps: int = 1,
+):
     """Save dict data to '*.vtu' file.
 
     Args:
         filename (str): Output filename.
-        data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Data in dict.
+        data_dict (Dict[str, np.ndarray]): Data in dict.
         coord_keys (Tuple[str, ...]): Tuple of coord key. such as ("x", "y").
         value_keys (Tuple[str, ...]): Tuple of value key. such as ("u", "v").
         num_timestamps (int, optional): Number of timestamp in data_dict. Defaults to 1.
@@ -109,17 +114,9 @@ def save_vtu_from_dict(filename, data_dict, coord_keys, value_keys, num_timestam
     coord = [data_dict[k] for k in coord_keys if k != "t"]
     value = [data_dict[k] for k in value_keys] if value_keys else None
 
-    if isinstance(coord[0], paddle.Tensor):
-        coord = [x.numpy() for x in coord]
-    else:
-        coord = [x for x in coord]
     coord = np.concatenate(coord, axis=1)
 
     if value is not None:
-        if isinstance(value[0], paddle.Tensor):
-            value = [x.numpy() for x in value]
-        else:
-            value = [x for x in value]
         value = np.concatenate(value, axis=1)
 
     _save_vtu_from_array(filename, coord, value, value_keys, num_timestamps)
@@ -128,26 +125,26 @@ def save_vtu_from_dict(filename, data_dict, coord_keys, value_keys, num_timestam
 def save_vtu_to_mesh(
     filename: str,
     data_dict: Dict[str, np.ndarray],
-    value_keys: Tuple[str, ...],
     coord_keys: Tuple[str, ...],
-    num_timestamps: int = 1,
+    value_keys: Tuple[str, ...],
 ):
     """Save data into .vtu format by meshio.
 
     Args:
         filename (str): File name.
-        data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Data in dict.
+        data_dict (Dict[str, np.ndarray]): Data in dict.
         coord_keys (Tuple[str, ...]): Tuple of coord key. such as ("x", "y").
         value_keys (Tuple[str, ...]): Tuple of value key. such as ("u", "v").
-        num_timestamp (int, optional): Number of timestamp in data_dict. Defaults to 1.
     """
-    path = os.path.dirname(filename)
-    os.makedirs(path, exist_ok=True)
+    npoint = len(next(iter(data_dict.values())))
+    coord_ndim = len(coord_keys)
 
-    n = len(next(iter(data_dict.values())))
-    m = len(coord_keys)
     # get the list variable transposed
-    points = np.stack((data_dict[key] for key in coord_keys)).reshape(m, n)
-    mesh = meshio.Mesh(points=points.T, cells=[("vertex", np.arange(n).reshape(n, 1))])
+    points = np.stack((data_dict[key] for key in coord_keys)).reshape(
+        coord_ndim, npoint
+    )
+    mesh = meshio.Mesh(
+        points=points.T, cells=[("vertex", np.arange(npoint).reshape(npoint, 1))]
+    )
     mesh.point_data = {key: data_dict[key] for key in value_keys}
     mesh.write(filename)