add save_tecplot_file

Author: HydrogenSulfate

docs/zh/api/utils/writer.md

@@ -5,5 +5,6 @@
     options:
       members:
         - save_csv_file
+        - save_tecplot_file
       show_root_heading: True
       heading_level: 3

ppsci/utils/__init__.py

@@ -35,6 +35,7 @@
 from ppsci.utils.save_load import save_checkpoint
 from ppsci.utils.symbolic import lambdify
 from ppsci.utils.writer import save_csv_file
+from ppsci.utils.writer import save_tecplot_file
 
 __all__ = [
     "AttrDict",
@@ -52,6 +53,7 @@
     "load_vtk_file",
     "load_vtk_with_time_file",
     "save_csv_file",
+    "save_tecplot_file",
     "dynamic_import_to_globals",
     "run_check",
     "run_check_mesh",

ppsci/utils/writer.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import csv
+import os
 from typing import Dict
 from typing import Optional
 from typing import Tuple
@@ -31,8 +32,8 @@
 
 
 def save_csv_file(
-    file_path: str,
-    data_dict: Dict[str, Union[np.ndarray, paddle.Tensor]],
+    filename: str,
+    data_dict: Dict[str, Union[np.ndarray, "paddle.Tensor"]],
     keys: Tuple[str, ...],
     alias_dict: Optional[Dict[str, str]] = None,
     use_header: bool = True,
@@ -42,11 +43,11 @@ def save_csv_file(
     """Write numpy data to csv file.
 
     Args:
-        file_path (str): Dump file path.
+        filename (str): Dump file path.
         data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Numpy data in dict.
         keys (Tuple[str, ...]): Keys for data_dict to be fetched.
         alias_dict (Optional[Dict[str, str]], optional): Alias dict for keys,
-            i.e. {dict_key: dump_key}. Defaults to None.
+            i.e. {dump_key: dict_key}. Defaults to None.
         use_header (bool, optional): Whether save csv with header. Defaults to True.
         delimiter (str, optional): Delemiter for splitting different data field. Defaults to ",".
         encoding (str, optional): Encoding. Defaults to "utf-8".
@@ -61,14 +62,14 @@ def save_csv_file(
         >>> save_csv_file(
         ...     "test.csv",
         ...     data_dict,
-        ...     ("a", "b"),
-        ...     alias_dict={"a": "A", "b": "B"},
+        ...     ("A", "B"),
+        ...     alias_dict={"A": "a", "B": "b"},
         ...     use_header=True,
         ...     delimiter=",",
         ...     encoding="utf-8",
         ... )
         >>> # == test.csv ==
-        >>> # a,b
+        >>> # A,B
         >>> # 1,4.12
         >>> # 2,5.25
         >>> # 3,6.337
@@ -81,28 +82,120 @@ def save_csv_file(
     data_fields = []
     header = []
     for key in keys:
-        if key not in data_dict:
-            raise KeyError(f"key({key}) do not exist in data_dict.")
-
-        data = data_dict[key]
+        fetch_key = alias_dict[key] if key in alias_dict else key
+        data = data_dict[fetch_key]
         if isinstance(data, paddle.Tensor):
             data = data.numpy()  # [num_of_samples, ]
 
         data = data.flatten()
         data_fields.append(data)
 
-        dump_key = alias_dict[key] if key in alias_dict else key
-        header.append(dump_key)
+        header.append(key)
 
     assert len(header) == len(data_fields)
 
-    data_fields = zip(*data_fields)
-    with open(file_path, "w", newline="", encoding=encoding) as file:
+    data_fields = zip(*data_fields)  # transpose col data to row data
+    with open(filename, "w", newline="", encoding=encoding) as file:
         writer = csv.writer(file, delimiter=delimiter)
 
         if use_header:
             writer.writerow(header)
 
         writer.writerows(data_fields)
 
-    logger.message(f"csv file has been dumped to {file_path}")
+    logger.message(f"csv file has been dumped to {filename}")
+
+
+def save_tecplot_file(
+    filename: str,
+    data_dict: Dict[str, Union[np.ndarray, "paddle.Tensor"]],
+    keys: Tuple[str, ...],
+    alias_dict: Optional[Dict[str, str]] = None,
+    delimiter: str = " ",
+    encoding: str = "utf-8",
+    num_timestamps: int = 1,
+):
+    """Write numpy data to tecplot file.
+
+    Args:
+        filename (str): Tecplot file path.
+        data_dict (Dict[str, Union[np.ndarray, paddle.Tensor]]): Numpy or Tensor data in dict.
+        keys (Tuple[str, ...]): Target keys to be dumped.
+        alias_dict (Optional[Dict[str, str]], optional): Alias dict for keys,
+            i.e. {dump_key: dict_key}. Defaults to None.
+        delimiter (str, optional): Delemiter for splitting different data field. Defaults to ",".
+        encoding (str, optional): Encoding. Defaults to "utf-8".
+        num_timestamps (int, optional): Number of timestamp over coord and value. Defaults to 1.
+
+    Examples:
+        >>> import numpy as np
+        >>> from ppsci.utils import save_tecplot_file
+        >>> data_dict = {
+        ...     "x": np.array([[1.0], [2.0], [10.0], [20.0], [100.0], [200.0]]), # [6, 1]
+        ...     "y": np.array([[-1.0], [-2.0], [-10.0], [-20.0], [-100.0], [-200.0]]) # [6, 1]
+        ... }
+        >>> save_tecplot_file(
+        ...    "./test.tec",
+        ...    data_dict,
+        ...    ("X", "Y"),
+        ...    alias_dict={"X": "x", "Y": "y"},
+        ...    num_timestamps=3,
+        ... )
+        >>> # == test_t-0.tec ==
+        >>> # title="./test_t-0.tec"
+        >>> # variables = "X", "Y"
+        >>> # Zone I = 2, J = 1, F = POINT
+        >>> # -1.0 1.0
+        >>> # -2.0 2.0
+
+        >>> # == test_t-1.tec ==
+        >>> # title="./test_t-1.tec"
+        >>> # variables = "X", "Y"
+        >>> # Zone I = 2, J = 1, F = POINT
+        >>> # -10.0 10.0
+        >>> # -20.0 20.0
+
+        >>> # == test_t-2.tec ==
+        >>> # title="./test_t-2.tec"
+        >>> # variables = "X", "Y"
+        >>> # Zone I = 2, J = 1, F = POINT
+        >>> # -100.0 100.0
+        >>> # -200.0 200.0
+    """
+    ntxy = len(next(iter(data_dict.values())))
+    if ntxy % num_timestamps != 0:
+        raise ValueError(
+            f"num_points({ntxy}) must be a multiple of "
+            f"num_timestamps({num_timestamps})."
+        )
+    nxy = ntxy // num_timestamps
+
+    os.makedirs(os.path.dirname(filename), exist_ok=True)
+
+    if filename.endswith(".tec"):
+        filename = filename[:-4]
+
+    for t in range(num_timestamps):
+        # write 1 tecplot file for each timestep
+        if num_timestamps > 1:
+            dump_filename = f"{filename}_t-{t}.tec"
+        else:
+            dump_filename = f"{filename}.tec"
+
+        with open(dump_filename, "w", encoding=encoding) as f:
+            # write meta information of tec
+            f.write(f'title="{dump_filename}"\n')
+            fetch_keys = {(alias_dict[key] if alias_dict else key) for key in keys}
+            header = ", ".join([f'"{key}"' for key in keys])
+            f.write(f"variables = {header}\n")
+            f.write(f"Zone I = {nxy}, J = 1, F = POINT\n")
+
+            # write points data into file
+            data_cur_time_step = [
+                data_dict[key][t * nxy : (t + 1) * nxy] for key in fetch_keys
+            ]
+
+            for items in zip(*data_cur_time_step):
+                f.write(delimiter.join([str(float(x)) for x in items]) + "\n")
+
+        logger.message(f"csv file has been dumped to {dump_filename}")