render rollout merge

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

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

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

gns/render_rollout.py

@@ -271,15 +271,21 @@ 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
 
                     # Create a color field based on material property and particle type
                     color_field = np.copy(material_property)
-                    static_particle_value = np.max(
-                        material_property) + 1  # Use a value outside the material property range
-                    color_field[particle_type == 3] = static_particle_value  # Assumes static particle type = 3
+                    static_particle_value = (
+                        np.max(material_property) + 1
+                    )  # Use a value outside the material property range
+                    color_field[particle_type == 3] = (
+                        static_particle_value  # Assumes static particle type = 3
+                    )
                     data["color"] = color_field
                 else:
                     # If no material property, use particle type for color
@@ -306,8 +312,10 @@ def write_vtk(self):
 
                 gridToVTK(
                     f"{path}/boundary{i}",
-                    x, y, z,
-                    cellData={"boundary": np.ones((1, 1, 1))}
+                    x,
+                    y,
+                    z,
+                    cellData={"boundary": np.ones((1, 1, 1))},
                 )
 
         print(f"vtk saved to: {self.output_dir}{self.output_name}...")

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

requirements.txt

@@ -15,3 +15,4 @@ torch_scatter
 torch-cluster
 tqdm
 toml
+pyvista