From 7190e4a599ac140b056517e50f753f6541ff5173 Mon Sep 17 00:00:00 2001
From: Naveen Raj Manoharan <150385340+Naveen-Raj-M@users.noreply.github.com>
Date: Tue, 22 Oct 2024 21:30:03 -0500
Subject: [PATCH] render rollout merge (#87)

* render rollout merge

* Update config.yaml

* Added unit test for render-rollout merge

* deleted debugging fixtures

* update config

* add test for VTK rendering

* bug fix for NoneType material property

* remove test_rendering and temp directory

* add test for vtk rendering

* modify config for render_rollout merge

* update to merge render-rollout

* set default mode to gif

* rewrite 'rendering' function in an extensible way

* improve readability and consistency

* update rendering options

* run black

* minor fix on viewpoint_rotation type

* improve logging and reformat with black

* refactor: move n_files function to a separate count_n_files.py in utils directory

* rename count_n_files.py to file_utils.py

* minor fix on module import

* run black

* minor fix on raising error

* add package for reading vtk files

---------

Co-authored-by: Naveen Raj Manoharan <naveen_raj_manoharan@login1.ls6.tacc.utexas.edu>
Co-authored-by: Naveen Raj Manoharan <naveen_raj_manoharan@v330-006.ls6.tacc.utexas.edu>
---
 README.md             |  21 ++-
 config.yaml           |  10 ++
 gns/args.py           |  15 +++
 gns/render_rollout.py |   5 +-
 gns/train.py          |  41 +++++-
 requirements.txt      |   1 +
 test/test_vtk.py      | 293 ++++++++++++++++++++++++++++++++++++++++++
 utils/file_utils.py   |  29 +++++
 8 files changed, 408 insertions(+), 7 deletions(-)
 create mode 100644 test/test_vtk.py
 create mode 100644 utils/file_utils.py

diff --git a/README.md b/README.md
index d1f70ac..07f8732 100644
--- a/README.md
+++ b/README.md
@@ -34,13 +34,19 @@ python3 -m gns.train mode="train" training.resume=True
 python3 -m meshnet.train mode="train" training.resume=True
 ```
 
-> Rollout prediction
+> Rollout prediction and render
 ```shell
 # For particulate domain,
 python3 -m gns.train mode="rollout"
 # For mesh-based domain,
 python3 -m meshnet.train mode="rollout"
 ```
+### Rendering Options
+Set rendering mode with `rendering.mode=<option>`
+`null`: Disables rendering
+`gif`: Creates a .gif file (default)
+`vtk`: Writes .vtu files for ParaView visualization
+Example: `rendering.mode=null` to disable rendering
 
 > Render
 ```shell
@@ -50,8 +56,6 @@ python3 -m gns.render_rollout --output_mode="gif" --rollout_dir="<path-containin
 python3 -m gns.render --rollout_dir="<path-containing-rollout-file>" --rollout_name="<name-of-rollout-file>"
 ```
 
-In particulate domain, the renderer also writes `.vtu` files to visualize in ParaView.
-
 ![Sand rollout](docs/img/rollout_0.gif)
 > GNS prediction of Sand rollout after training for 2 million steps.
 
@@ -118,8 +122,16 @@ hardware:
 # Logging configuration
 logging:
   tensorboard_dir: logs/
-```
 
+# Rendering configuration
+rendering:
+  mode: gif
+  
+  gif:
+    step_stride: 3
+    vertical_camera_angle: 20
+    viewpoint_rotation: 0.3
+    change_yz: False
 </details>
 
 
@@ -180,7 +192,6 @@ The total number of training steps to execute before stopping.
 **nsave_steps (Integer)**
 
 Interval at which the model and training state are saved.
-
 </details>
 
 ## Datasets
diff --git a/config.yaml b/config.yaml
index fdd3229..8f1e9ba 100644
--- a/config.yaml
+++ b/config.yaml
@@ -49,3 +49,13 @@ hardware:
 # Logging configuration
 logging:
   tensorboard_dir: logs/
+
+# Rendering configuration
+rendering:
+  mode: gif
+  
+  gif:
+    step_stride: 3
+    vertical_camera_angle: 20
+    viewpoint_rotation: 0.3
+    change_yz: False
\ No newline at end of file
diff --git a/gns/args.py b/gns/args.py
index dd3dcd6..2fbe4f8 100644
--- a/gns/args.py
+++ b/gns/args.py
@@ -53,6 +53,20 @@ class LoggingConfig:
     tensorboard_dir: str = "logs/"
 
 
+@dataclass
+class GifConfig:
+    step_stride: int = 3
+    vertical_camera_angle: int = 20
+    viewpoint_rotation: float = 0.3
+    change_yz: bool = False
+
+
+@dataclass
+class RenderingConfig:
+    mode: Optional[str] = field(default="gif")
+    gif: GifConfig = field(default_factory=GifConfig)
+
+
 @dataclass
 class Config:
     mode: str = "train"
@@ -62,6 +76,7 @@ class Config:
     training: TrainingConfig = field(default_factory=TrainingConfig)
     hardware: HardwareConfig = field(default_factory=HardwareConfig)
     logging: LoggingConfig = field(default_factory=LoggingConfig)
+    rendering: RenderingConfig = field(default_factory=RenderingConfig)
 
 
 # Hydra configuration
diff --git a/gns/render_rollout.py b/gns/render_rollout.py
index fcf175a..dcac88f 100644
--- a/gns/render_rollout.py
+++ b/gns/render_rollout.py
@@ -271,7 +271,10 @@ def write_vtk(self):
                 }
 
                 # Check if material property exists and add it to data if it does
-                if "material_property" in self.rollout_data:
+                if (
+                    "material_property" in self.rollout_data
+                    and self.rollout_data["material_property"] is not None
+                ):
                     material_property = self.rollout_data["material_property"]
                     data["material_property"] = material_property
 
diff --git a/gns/train.py b/gns/train.py
index 8d8ed6a..a730e75 100644
--- a/gns/train.py
+++ b/gns/train.py
@@ -21,6 +21,7 @@
 from gns import reading_utils
 from gns import particle_data_loader as pdl
 from gns import distribute
+from gns import render_rollout
 from gns.args import Config
 
 Stats = collections.namedtuple("Stats", ["mean", "std"])
@@ -200,15 +201,53 @@ def predict(device: str, cfg: DictConfig):
                 example_rollout["metadata"] = metadata
                 example_rollout["loss"] = loss.mean()
                 filename = f"{cfg.output.filename}_ex{example_i}.pkl"
-                filename = os.path.join(cfg.output.path, filename)
+                filename_render = f"{cfg.output.filename}_ex{example_i}"
+                filename = os.path.join(cfg.output.path, filename_render)
                 with open(filename, "wb") as f:
                     pickle.dump(example_rollout, f)
+            if cfg.rendering.mode:
+                rendering(cfg.output.path, filename_render, cfg)
 
     print(
         "Mean loss on rollout prediction: {}".format(torch.mean(torch.cat(eval_loss)))
     )
 
 
+def rendering(input_dir, input_name, cfg: DictConfig):
+    """
+    Render output based on the specified configuration and input parameters.
+    It supports rendering in both GIF and VTK formats.
+
+    Args:
+        input_dir (str): The directory containing the input files for rendering.
+        input_name (str): The base name of the input file to be rendered.
+        cfg (DictConfig): The configuration dictionary that specifies rendering options,
+                          including the rendering mode and relevant parameters for that mode.
+
+    Raises:
+        ValueError: If the specified rendering mode is not supported.
+    """
+    render = render_rollout.Render(input_dir, input_name)
+
+    rendering_modes = {
+        "gif": lambda: render.render_gif_animation(
+            point_size=1,
+            timestep_stride=cfg.rendering.gif.step_stride,
+            vertical_camera_angle=cfg.rendering.gif.vertical_camera_angle,
+            viewpoint_rotation=cfg.rendering.gif.viewpoint_rotation,
+            change_yz=cfg.rendering.gif.change_yz,
+        ),
+        "vtk": lambda: render.write_vtk(),
+    }
+
+    if cfg.rendering.mode in ["gif", "vtk"]:
+        rendering_mode = rendering_modes.get(cfg.rendering.mode)
+        rendering_mode()
+
+    else:
+        raise ValueError(f"Unsupported rendering mode: {cfg.rendering.mode}")
+
+
 def optimizer_to(optim, device):
     for param in optim.state.values():
         # Not sure there are any global tensors in the state dict
diff --git a/requirements.txt b/requirements.txt
index 595fcd9..da9a0a2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -15,3 +15,4 @@ torch_scatter
 torch-cluster
 tqdm
 toml
+pyvista
\ No newline at end of file
diff --git a/test/test_vtk.py b/test/test_vtk.py
new file mode 100644
index 0000000..7079e2d
--- /dev/null
+++ b/test/test_vtk.py
@@ -0,0 +1,293 @@
+import pytest
+import numpy as np
+import pickle
+import os
+import tempfile
+import shutil
+import pyvista as pv
+import logging
+from gns.train import rendering
+from omegaconf import DictConfig
+from typing import Tuple, Dict
+from utils.file_utils import n_files
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@pytest.fixture
+def cfg_vtk() -> DictConfig:
+    """
+    Fixture for VTK configuration.
+
+    Returns:
+        DictConfig: Configuration dictionary for VTK rendering mode.
+    """
+    logger.info("Setting up VTK configuration.")
+    return DictConfig({"rendering": {"mode": "vtk"}})
+
+
+@pytest.fixture
+def dummy_pkl_data() -> dict:
+    """
+    Fixture for generating dummy pickle data for testing.
+
+    Returns:
+        Dict: A dictionary containing dummy rollout data.
+    """
+    n_timesteps = 2
+    n_particles = 3
+    dim = 2
+    n_init_pos = 2
+
+    try:
+        # Generate random predictions and ground truth positions
+        predictions = np.random.rand(n_timesteps, n_particles, dim)
+        ground_truth_positions = np.random.randn(n_timesteps, n_particles, dim)
+        loss = (predictions - ground_truth_positions) ** 2
+
+        # Rollout dictionary to store all relevant information
+        dummy_rollout = {
+            "initial_positions": np.random.rand(n_init_pos, n_particles, dim),
+            "predicted_rollout": predictions,
+            "ground_truth_rollout": ground_truth_positions,
+            "particle_types": np.full(n_particles, 5),
+            "metadata": {"bounds": [[0.0, 1.0], [0.0, 1.0]]},
+            # MSE loss between predictions and ground truth positions
+            "loss": loss.mean(),
+        }
+        logger.info(
+            f"Generated dummy data: {n_particles} particles over {n_timesteps} time steps in {dim} dimensions, "
+            f"with {n_init_pos} initial positions"
+        )
+
+    except Exception as e:
+        logger.error(f"Failed to generate dummy pickle data: {e}")
+        raise
+
+    return dummy_rollout
+
+
+@pytest.fixture
+def temp_dir_with_file(dummy_pkl_data: Dict) -> Tuple[str, str]:
+    """
+    Fixture to create a temporary directory and a pickle file containing the
+    provided dummy data for testing purposes.
+
+    Args:
+        dummy_pkl_data (dict): A dictionary containing the dummy data to be
+        serialized and stored in a temporary pickle file.
+
+    Yields:
+        Tuple[str, str]: A tuple containing:
+            - The path to the temporary directory where the pickle file is stored.
+            - The base name of the pickle file (without the '.pkl' extension).
+    """
+    temp_dir = tempfile.mkdtemp()
+    logger.info(f"Created temporary directory: {temp_dir}")
+
+    try:
+        with tempfile.NamedTemporaryFile(
+            dir=temp_dir, suffix=".pkl", delete=False
+        ) as temp_file:
+            pkl_file_path = temp_file.name
+
+            with open(pkl_file_path, "wb") as f:
+                pickle.dump(dummy_pkl_data, f)
+
+            # get the base file name without any extension
+            file_name = os.path.splitext(os.path.basename(pkl_file_path))[0]
+            logger.info(f"created temporary '.pkl' file: {file_name}")
+            temp_dir += "/"
+
+            yield temp_dir, file_name
+
+    except Exception as e:
+        logger.error(
+            f"Failed to create a Temporary file for the input rollout data: {e}"
+        )
+        raise
+
+    finally:
+        shutil.rmtree(temp_dir)
+
+
+def verify_vtk_files(rollout_data: dict, label: str, temp_dir: str) -> None:
+    """
+    Verify the integrity of VTK files against expected data.
+
+    This function checks VTK files (VTU and VTR) for specific properties, ensuring that the
+    displacement, particle types, color maps, and bounds match the expected values in the
+    provided rollout data.
+
+    Args:
+        rollout_data (dict): A dictionary containing
+            - 'inital positions'
+            - 'predicted_rollout'
+            - 'ground_truth_rollout'
+            - 'particle_types'
+            - 'metadata'
+        label (str): The label for the specific rollout data to verify against.
+        temp_dir (str): The temporary directory where the VTK files are stored.
+
+    Raises:
+        AssertionError: If any of the checks on displacement, particle types, color maps,
+        or bounds fail.
+        Exception: If there is an error reading or processing the VTK files.
+    """
+    VTU_PREFIX = "points"
+    VTR_PREFIX = "boundary"
+    VTU_EXTENSION = "vtu"
+    VTR_EXTENSION = "vtr"
+
+    positions = np.concatenate(
+        [rollout_data["initial_positions"], rollout_data[label]],
+        axis=0,
+    )
+
+    for time_step in range(positions.shape[0]):
+        try:
+            vtu_path = os.path.join(
+                temp_dir, f"{VTU_PREFIX}{time_step}.{VTU_EXTENSION}"
+            )
+            logger.info(f"Reading VTU file: {vtu_path}")
+            vtu_object = pv.read(vtu_path)
+
+            displacement = vtu_object["displacement"]
+            particle_type = vtu_object["particle_type"]
+            color_map = vtu_object["color"]
+
+            assert np.all(
+                displacement
+                == np.linalg.norm(positions[0] - positions[time_step], axis=1)
+            ), (
+                f"Displacement mismatch for timestep {time_step}: "
+                f"expected {np.linalg.norm(positions[0] - positions[time_step], axis=1)}, "
+                f"got {displacement}"
+            )
+            assert np.all(particle_type == rollout_data["particle_types"]), (
+                f"Particle type mismatch for timestep {time_step}: "
+                f"expected {rollout_data['particle_types']}, got {particle_type}"
+            )
+            assert np.all(color_map == rollout_data["particle_types"]), (
+                f"Color map mismatch for timestep {time_step}: "
+                f"expected {rollout_data['particle_types']}, got {color_map}"
+            )
+
+        except AssertionError as e:
+            logger.error(
+                f"Assertion failed while verifying {VTU_PREFIX}{time_step}.{VTU_EXTENSION}: {e}"
+            )
+            raise
+        except Exception as e:
+            logger.error(f"Error reading {VTU_PREFIX}{time_step}.{VTU_EXTENSION}: {e}")
+            raise
+
+        try:
+            vtr_path = os.path.join(
+                temp_dir, f"{VTR_PREFIX}{time_step}.{VTR_EXTENSION}"
+            )
+            logger.info(f"Reading VTR file: {vtr_path}")
+            vtr_object = pv.read(vtr_path)
+
+            bounds = vtr_object.bounds
+            xmin, xmax, ymin, ymax, zmin, zmax = bounds
+
+            assert xmin == rollout_data["metadata"]["bounds"][0][0], (
+                f"Xmin mismatch for timestep {time_step}: "
+                f"expected {rollout_data['metadata']['bounds'][0][0]}, got {xmin}"
+            )
+            assert xmax == rollout_data["metadata"]["bounds"][0][1], (
+                f"Xmax mismatch for timestep {time_step}: "
+                f"expected {rollout_data['metadata']['bounds'][0][1]}, got {xmax}"
+            )
+            assert ymin == rollout_data["metadata"]["bounds"][1][0], (
+                f"Ymin mismatch for timestep {time_step}: "
+                f"expected {rollout_data['metadata']['bounds'][1][0]}, got {ymin}"
+            )
+            assert ymax == rollout_data["metadata"]["bounds"][1][1], (
+                f"Ymax mismatch for timestep {time_step}: "
+                f"expected {rollout_data['metadata']['bounds'][1][1]}, got {ymax}"
+            )
+
+        except AssertionError as e:
+            logger.error(
+                f"Assertion failed while verifying {VTR_PREFIX}{time_step}.{VTR_EXTENSION}: {e}"
+            )
+            raise
+        except Exception as e:
+            logger.error(f"Error reading {VTR_PREFIX}{time_step}.{VTR_EXTENSION}: {e}")
+            raise
+
+    logger.info("VTK file verification complete.")
+
+
+def test_rendering_vtk(
+    temp_dir_with_file: Tuple[str, str], cfg_vtk: DictConfig
+) -> None:
+    """
+    Test the VTK rendering function.
+
+    Args:
+        temp_dir_with_file (Tuple[str, str]): Tuple containing the path to the temporary directory and the pickle file name.
+        cfg_vtk (DictConfig): Configuration dictionary for VTK rendering mode.
+    """
+    input_dir, input_file = temp_dir_with_file
+
+    logger.info(f"Starting rendering test: {input_file}")
+
+    try:
+        rendering(input_dir, input_file, cfg_vtk)
+    except Exception as e:
+        logger.error("Failed to render with dummy rollout data")
+        raise
+
+    try:
+        # Define paths for the generated VTK files
+        VTK_GNS_SUFFIX = "_vtk-GNS"
+        VTK_REALITY_SUFFIX = "_vtk-Reality"
+
+        vtk_path_gns = os.path.join(input_dir, f"{input_file}{VTK_GNS_SUFFIX}")
+        vtk_path_reality = os.path.join(input_dir, f"{input_file}{VTK_REALITY_SUFFIX}")
+
+        logger.info("Loading rollout data from pickle file.")
+        with open(os.path.join(input_dir, f"{input_file}.pkl"), "rb") as file:
+            rollout = pickle.load(file)
+
+        # Count the number of .vtu and .vtr files in the VTK directories
+        n_vtu_files_gns = n_files(vtk_path_gns, "vtu")
+        n_vtu_files_reality = n_files(vtk_path_reality, "vtu")
+        n_vtr_files_gns = n_files(vtk_path_gns, "vtr")
+        n_vtr_files_reality = n_files(vtk_path_reality, "vtr")
+
+        expected_n_files = (
+            rollout["initial_positions"].shape[0]
+            + rollout["predicted_rollout"].shape[0]
+        )
+        logger.info(f"Expected number of files: {expected_n_files}.")
+
+        # Assert that the number of .vtu and .vtr files matches the expected count
+        assert (
+            n_vtu_files_gns == expected_n_files
+        ), f"Expected {expected_n_files} VTU files in GNS path, got {n_vtu_files_gns}"
+        assert (
+            n_vtu_files_reality == expected_n_files
+        ), f"Expected {expected_n_files} VTU files in Reality path, got {n_vtu_files_reality}"
+        assert (
+            n_vtr_files_gns == expected_n_files
+        ), f"Expected {expected_n_files} VTR files in GNS path, got {n_vtr_files_gns}"
+        assert (
+            n_vtr_files_reality == expected_n_files
+        ), f"Expected {expected_n_files} VTR files in Reality path, got {n_vtr_files_reality}"
+
+        logger.info("Verifying VTK files for predicted rollout.")
+        verify_vtk_files(rollout, "predicted_rollout", vtk_path_gns)
+        logger.info("Verifying VTK files for ground truth rollout.")
+        verify_vtk_files(rollout, "ground_truth_rollout", vtk_path_reality)
+
+        logger.info("Rendering test completed successfully.")
+
+    except Exception as e:
+        logger.error(f"Rendering test failed: {e}")
+        raise
diff --git a/utils/file_utils.py b/utils/file_utils.py
new file mode 100644
index 0000000..5fed43c
--- /dev/null
+++ b/utils/file_utils.py
@@ -0,0 +1,29 @@
+import os
+import glob
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def n_files(directory: str, extension: str) -> int:
+    """
+    Count the number of files with a specific extension in a directory.
+
+    Args:
+        directory (str): Directory path.
+        extension (str): File extension to count.
+
+    Returns:
+        int: Number of files with the specified extension.
+    """
+    try:
+        pattern = os.path.join(directory, f"*.{extension}")
+        file_count = len(glob.glob(pattern))
+        return file_count
+    except Exception as e:
+        logger.error(
+            f"Error counting files in {directory} with extension '{extension}': {e}"
+        )
+        return 0