From fa2caaee8c037986179e41e8c626f0582d0eb234 Mon Sep 17 00:00:00 2001 From: Vincent Lannelongue Date: Mon, 23 Mar 2026 18:42:53 +0100 Subject: [PATCH 1/6] Add xdmf append function and update lightning module --- graphphysics/training/lightning_module.py | 138 +++++++++------------- graphphysics/utils/meshio_mesh.py | 122 ++++++++++++++++++- 2 files changed, 174 insertions(+), 86 deletions(-) diff --git a/graphphysics/training/lightning_module.py b/graphphysics/training/lightning_module.py index 9ad28cf..7340aa6 100644 --- a/graphphysics/training/lightning_module.py +++ b/graphphysics/training/lightning_module.py @@ -1,9 +1,7 @@ import os -import shutil from typing import Dict, List, Optional import lightning as L -import meshio import torch import torch.nn as nn from loguru import logger @@ -17,7 +15,11 @@ get_simulator, ) from graphphysics.utils.loss import L2Loss, MultiLoss -from graphphysics.utils.meshio_mesh import convert_to_meshio_vtu +from graphphysics.utils.meshio_mesh import ( + append_mesh_to_xdmf, + convert_to_meshio_vtu, + meshes_to_xdmf, +) from graphphysics.utils.nodetype import NodeType from graphphysics.utils.scheduler import CosineWarmupScheduler @@ -108,13 +110,9 @@ def __init__( self.previous_data_start = previous_data_start self.previous_data_end = previous_data_end - # For one trajectory vizualization - self.trajectory_to_save: list[Batch] = [] - # Prediction self.prediction_save_path: str = prediction_save_path self.current_pred_trajectory = 0 - self.prediction_trajectory: list[Batch] = [] self.last_pred_prediction = None self.last_previous_data_pred_prediction = None @@ -148,6 +146,9 @@ def __init__( if self.use_spatial_mtp: self._setup_spatial_mtp(processor, device) + # TODO: Decide on whether or not to keep this + # self.compress_predictions = parameters["compression"] + def forward(self, graph: Batch): return self.model(graph) @@ -334,42 +335,27 @@ def teardown(self, stage: Optional[str] = None) -> None: self._remove_spatial_mtp_hooks() super().teardown(stage) - def _save_trajectory_to_xdmf( - self, - trajectory: list[Batch], - save_dir: str, - archive_filename: str, - timestep: float = 1, + def _save_batch_to_xdmf( + self, batch: Batch, save_dir: str, archive_filename: str, timestep: float = 1 ): + """ + Saves a batch to an XDMF/H5 archive. + Creates the archive if it doesn't exist, appends the batch at the end it does. + """ os.makedirs(save_dir, exist_ok=True) archive_path = os.path.join(save_dir, archive_filename) - xdmf_filename = f"{archive_path}.xdmf" - init_mesh = convert_to_meshio_vtu(trajectory[0], add_all_data=True) - points = init_mesh.points - cells = init_mesh.cells - try: - with meshio.xdmf.TimeSeriesWriter(xdmf_filename) as writer: - # Write the mesh (points and cells) once - writer.write_points_cells(points, cells) - # Loop through time steps and write data - t = timestep if not self.use_previous_data else 2 * timestep - for idx, graph in enumerate(trajectory): - mesh = convert_to_meshio_vtu(graph, add_all_data=True) - point_data = mesh.point_data - cell_data = mesh.cell_data - writer.write_data(t, point_data=point_data, cell_data=cell_data) - t += timestep - - except Exception as e: - logger.error(f"Error saving graph {idx} at epoch {self.current_epoch}: {e}") - logger.info( - f"Validation Trajectory {archive_filename.split('_')[-1]} saved at {save_dir}." - ) - h5_filename = xdmf_filename.replace(".xdmf", ".h5") - src = os.path.join(os.getcwd(), os.path.basename(h5_filename)) - # The h5 file may be in the cwd (meshio bug), move it to xdmf location - if os.path.exists(src): - shutil.move(src=src, dst=h5_filename) + mesh = convert_to_meshio_vtu(batch, add_all_data=True) + if not os.path.exists(f"{archive_path}.h5") or not os.path.exists( + f"{archive_path}.xdmf" + ): + meshes_to_xdmf(filename=archive_path, meshes=[mesh], timestep=timestep) + else: + append_mesh_to_xdmf( + filename=archive_path, + mesh=mesh, + timestep=timestep, + compress=self.compress_predictions, + ) def _reset_validation_trajectory(self): self.current_val_trajectory += 1 @@ -430,7 +416,16 @@ def validation_step(self, batch: Batch, batch_idx: int): ) if self.current_val_trajectory == 0: - self.trajectory_to_save.append(batch) + self._save_batch_to_xdmf( + batch, + os.path.join("meshes", f"epoch_{self.current_epoch}"), + self._get_frame_savename( + batch, + self.current_val_trajectory, + prefix=f"graph_epoch_{self.current_epoch}", + ), + timestep=self.timestep, + ) node_type = batch.x[:, self.model.node_type_index] self.val_step_outputs.append(predicted_outputs.cpu()) @@ -456,7 +451,6 @@ def _reset_validation_epoch_end(self): self.current_val_trajectory = 0 self.last_val_prediction = None self.last_previous_data_prediction = None - self.trajectory_to_save.clear() self.step_counter = 0 self.first_step_losses = [] @@ -482,20 +476,7 @@ def on_validation_epoch_end(self): mean_first_step_loss = torch.stack(self.first_step_losses).mean().item() self.log( "val_1step_rmse", mean_first_step_loss, on_epoch=True, prog_bar=True - ) - - # Save trajectory graphs - save_dir = os.path.join("meshes", f"epoch_{self.current_epoch}") - self._save_trajectory_to_xdmf( - self.trajectory_to_save, - save_dir, - self._get_traj_savename( - self.trajectory_to_save, - self.current_val_trajectory, - prefix=f"graph_epoch_{self.current_epoch}", - ), - timestep=self.timestep, - ) + ) # Clear stored outputs self._reset_validation_epoch_end() @@ -521,7 +502,6 @@ def configure_optimizers(self): def _reset_prediction_trajectory(self): self.current_pred_trajectory += 1 - self.prediction_trajectory = [] self.last_pred_prediction = None self.last_previous_data_pred_prediction = None @@ -533,16 +513,7 @@ def predict_step(self, batch: Batch): """ batch = batch.to(self.device, non_blocking=True) if batch.traj_index > self.current_pred_trajectory: - # save - self._save_trajectory_to_xdmf( - self.prediction_trajectory, - self.prediction_save_path, - self._get_traj_savename( - self.prediction_trajectory, self.current_pred_trajectory - ), - timestep=self.timestep, - ) - # reset + # reset when changing trajectory self._reset_prediction_trajectory() # predict @@ -555,28 +526,25 @@ def predict_step(self, batch: Batch): ) = self._make_prediction( batch, self.last_pred_prediction, self.last_previous_data_pred_prediction ) - self.prediction_trajectory.append(batch) + self._save_batch_to_xdmf( + batch, + self.prediction_save_path, + self._get_frame_savename( + batch, + self.current_pred_trajectory, + ), + timestep=self.timestep, + ) def _reset_predict_epoch_end(self): - self.prediction_trajectory.clear() self.last_pred_prediction = None self.last_previous_data_pred_prediction = None self.current_pred_trajectory = 0 def on_predict_epoch_end(self): """ - Save last trajectory to xdmf and clear stored outputs. + Clear stored outputs. """ - self._save_trajectory_to_xdmf( - self.prediction_trajectory, - self.prediction_save_path, - self._get_traj_savename( - self.prediction_trajectory, self.current_pred_trajectory - ), - timestep=self.timestep, - ) - - # Clear stored outputs self._reset_predict_epoch_end() def on_save_checkpoint(self, checkpoint: dict): @@ -594,19 +562,19 @@ def on_load_checkpoint(self, checkpoint): """ self.wandb_run_id = checkpoint.get("wandb_run_id", None) - def _get_traj_savename( - self, traj: list[Batch], traj_idx: int, prefix: str = "graph" + def _get_frame_savename( + self, batch: Batch, traj_idx: int, prefix: str = "graph" ) -> str: """ Get the name of the trajectory to save (id if provided in attributes, index otherwise). Args: - traj (list[Batch]): List of Batch objects representing the trajectory. + batch (Batch): Batch object. traj_idx (int): Index of the current trajectory. prefix (str): Prefix for the trajectory filename. (does not include trailing '_') Returns: str: The name of the trajectory to save (no extensions). """ - if hasattr(traj[0], "id") and traj[0].id[0] is not None: - return f"{prefix}_{traj[0].id[0]}" + if hasattr(batch, "id") and batch.id[0] is not None: + return f"{prefix}_{batch.id[0]}" else: return f"{prefix}_{traj_idx}" diff --git a/graphphysics/utils/meshio_mesh.py b/graphphysics/utils/meshio_mesh.py index c82557b..5a97e83 100644 --- a/graphphysics/utils/meshio_mesh.py +++ b/graphphysics/utils/meshio_mesh.py @@ -1,9 +1,12 @@ +import copy import os import shutil from typing import List +import h5py import meshio import numpy as np +from lxml import etree from torch_geometric.data import Data @@ -70,7 +73,7 @@ def vtu_to_xdmf( filename: str, files_list: List[str], timestep=1, remove_vtus: bool = True ) -> None: """ - Writes a time series of meshes (same points and cells) into XDMF/HDF5 format. + Writes a time series of meshes (same points and cells) into XDMF/HDF5 format from VTU files. Args: filename (str): Name for the XDMF/HDF5 file without the extension. @@ -111,3 +114,120 @@ def vtu_to_xdmf( if remove_vtus: for file in files_list: os.remove(file) + + +def meshes_to_xdmf( + filename: str, + meshes: List[meshio.Mesh], + timestep=1, +) -> None: + """ + Writes a time series of meshes (same points and cells) into XDMF/HDF5 format from meshio.Mesh objects. + + Args: + filename (str): Name for the XDMF/HDF5 file without the extension. + meshes (List[meshio.Mesh]): List of the meshes to compress. + timestep (float, optional): Timestep between to consecutive timeframes. + + Returns: + None: XDMF/HDF5 file is saved to the path filename. + """ + + h5_filename = f"{filename}.h5" + xdmf_filename = f"{filename}.xdmf" + + points = meshes[0].points + cells = meshes[0].cells + + # Open the TimeSeriesWriter for HDF5 + with meshio.xdmf.TimeSeriesWriter(xdmf_filename) as writer: + # Write the mesh (points and cells) once + writer.write_points_cells(points, cells) + + # Loop through time steps and write data + t = 0 + for mesh in meshes: + point_data = mesh.point_data + cell_data = mesh.cell_data + writer.write_data(t, point_data=point_data, cell_data=cell_data) + t += timestep + + # The H5 archive is systematically created in cwd, we just need to move it + shutil.move( + src=os.path.join(os.getcwd(), os.path.split(h5_filename)[1]), dst=h5_filename + ) + + +def append_mesh_to_xdmf( + filename: str, mesh: meshio.Mesh, timestep: float = 1.0, compress=False +) -> None: + """ + Appends a single timeframe to an existing XDMF/HDF5 time series archive, + without loading the existing data into RAM. + + Args: + filename (str): Path to the existing archive (without extension). + mesh (meshio.Mesh): Mesh object for the new timeframe. + Must share the same points and cells as the existing archive. + timestep (float): Time increment added to the last recorded timestep. + + Returns: + None: XDMF/HDF5 file is saved to the path filename. + """ + + h5_filename = f"{filename}.h5" + xdmf_filename = f"{filename}.xdmf" + + if not os.path.exists(h5_filename) or not os.path.exists(xdmf_filename): + raise FileNotFoundError(f"XDMF/HDF5 file not found: {filename}") + + # Get temporal grid to get the last timestep grid + tree = etree.parse(xdmf_filename) + root = tree.getroot() + temporal_grid = root.find( + ".//{*}Grid[@GridType='Collection'][@CollectionType='Temporal']" + ) + if temporal_grid is None: + raise ValueError( + "Could not find the temporal grid collection in the XDMF file." + ) + + time_grids = temporal_grid.findall("Grid") + + last_grid = time_grids[-1] + last_time = float(last_grid.find("Time").get("Value")) + new_time = last_time + timestep + + # Add data to the H5 file. + with h5py.File(h5_filename, "a") as h5: + existing_h5_keys = h5.keys() + last_h5_key_idx = max(int(k.replace("data", "")) for k in existing_h5_keys) + h5_new_keys_mapping = {} + for i, (field_name, field_values) in enumerate(mesh.point_data.items()): + new_key = f"data{last_h5_key_idx + (i+1)}" + if compress: + h5.create_dataset( + new_key, + data=np.asarray(field_values), + chunks=True, + compression="gzip", + compression_opts=4, # Default value for meshio, ranges from 0 to 9 + ) + else: + h5.create_dataset(new_key, data=np.asarray(field_values)) + h5_new_keys_mapping[field_name] = new_key + + # Add a new grid to the XDMF file. + new_grid = copy.deepcopy(last_grid) + new_grid.find("Time").set("Value", str(new_time)) + + for data_item in new_grid.iter(): + if data_item.tag == "Attribute": + next_h5_key = h5_new_keys_mapping[data_item.attrib["Name"]] + if data_item.tag == "DataItem": + data_item.text = f"{os.path.basename(h5_filename)}:/{next_h5_key}" + + temporal_grid.append(new_grid) + tree.write( + xdmf_filename, pretty_print=False, xml_declaration=False, encoding="UTF-8" + ) From 58d3cf6c6e58a28648a53da3f53de7f8d091eaf6 Mon Sep 17 00:00:00 2001 From: Vincent Lannelongue Date: Tue, 24 Mar 2026 11:21:06 +0100 Subject: [PATCH 2/6] Add tests for lightining module and meshiomesh and remove compress parameter --- graphphysics/training/lightning_module.py | 9 +- requirements.txt | 3 +- .../training/test_lightningmodule.py | 218 +++++++++--------- tests/graphphysics/utils/test_meshio_mesh.py | 178 +++++++++++++- 4 files changed, 294 insertions(+), 114 deletions(-) diff --git a/graphphysics/training/lightning_module.py b/graphphysics/training/lightning_module.py index 7340aa6..50001aa 100644 --- a/graphphysics/training/lightning_module.py +++ b/graphphysics/training/lightning_module.py @@ -350,12 +350,7 @@ def _save_batch_to_xdmf( ): meshes_to_xdmf(filename=archive_path, meshes=[mesh], timestep=timestep) else: - append_mesh_to_xdmf( - filename=archive_path, - mesh=mesh, - timestep=timestep, - compress=self.compress_predictions, - ) + append_mesh_to_xdmf(filename=archive_path, mesh=mesh, timestep=timestep) def _reset_validation_trajectory(self): self.current_val_trajectory += 1 @@ -476,7 +471,7 @@ def on_validation_epoch_end(self): mean_first_step_loss = torch.stack(self.first_step_losses).mean().item() self.log( "val_1step_rmse", mean_first_step_loss, on_epoch=True, prog_bar=True - ) + ) # Clear stored outputs self._reset_validation_epoch_end() diff --git a/requirements.txt b/requirements.txt index bcd56d3..7f6060d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,4 +17,5 @@ wandb==0.19.8 absl-py==2.2.0 wandb[media]==0.19.8 panel==1.6.1 -einops==0.8.1 \ No newline at end of file +einops==0.8.1 +lxml \ No newline at end of file diff --git a/tests/graphphysics/training/test_lightningmodule.py b/tests/graphphysics/training/test_lightningmodule.py index 8833326..1c9b90e 100644 --- a/tests/graphphysics/training/test_lightningmodule.py +++ b/tests/graphphysics/training/test_lightningmodule.py @@ -163,6 +163,25 @@ def test_validation_step(self): # Check that last_previous_data_prediction is not set self.assertIsNone(self.model.last_previous_data_prediction) + # Check that .xdmf is created after validation step + xdmf_path = os.path.join( + "meshes", + f"epoch_{self.model.current_epoch}", + f"graph_epoch_{self.model.current_epoch}_{self.model.current_val_trajectory}.xdmf", + ) + h5_path = os.path.join( + "meshes", + f"epoch_{self.model.current_epoch}", + f"graph_epoch_{self.model.current_epoch}_{self.model.current_val_trajectory}.h5", + ) + + self.assertTrue(os.path.exists(xdmf_path)) + self.assertTrue(os.path.exists(h5_path)) + + # Delete created validation files + os.remove(xdmf_path) + os.remove(h5_path) + def test_validation_step_w_previous_data(self): self.dataloader = DataLoader(self.dataset, batch_size=1) batch = next(iter(self.dataloader)) @@ -195,6 +214,25 @@ def test_validation_step_w_previous_data(self): self.model.previous_data_start = None self.model.previous_data_end = None + # Check that .xdmf is created after validation step + xdmf_path = os.path.join( + "meshes", + f"epoch_{self.model.current_epoch}", + f"graph_epoch_{self.model.current_epoch}_{self.model.current_val_trajectory}.xdmf", + ) + h5_path = os.path.join( + "meshes", + f"epoch_{self.model.current_epoch}", + f"graph_epoch_{self.model.current_epoch}_{self.model.current_val_trajectory}.h5", + ) + + self.assertTrue(os.path.exists(xdmf_path)) + self.assertTrue(os.path.exists(h5_path)) + + # Delete created validation files + os.remove(xdmf_path) + os.remove(h5_path) + def test_spatial_mtp_setup(self): params = deepcopy(self.parameters) params["training"] = {"use_spatial_mtp": True} @@ -241,30 +279,6 @@ def test_on_validation_epoch_end(self): output_dim = 2 self.model.eval() - # Simulate trajectory_to_save with sample graphs - num_graphs = 3 - for i in range(num_graphs): - # Create a simple graph - pos = torch.tensor( - [[0.0 + i, 0.0], [1.0 + i, 0.0], [1.0 + i, 1.0], [0.0 + i, 1.0]], - dtype=torch.float, - ) - edge_index = torch.tensor( - [[0, 1, 2, 3], [1, 2, 3, 0]], dtype=torch.long - ) - x = torch.tensor( - [ - [i * 10 + 1, i * 10 + 1], - [i * 10 + 2, i * 10 + 2], - [i * 10 + 3, i * 10 + 3], - [i * 10 + 4, i * 10 + 4], - ], - dtype=torch.float, - ) - face = torch.tensor([[0], [1], [2]]) - graph = Data(pos=pos, edge_index=edge_index, x=x, face=face) - self.model.trajectory_to_save.append(graph) - # Simulate val_step_outputs and val_step_targets for i in range(num_steps): predicted_outputs = torch.randn(batch_size, output_dim) @@ -291,23 +305,6 @@ def test_on_validation_epoch_end(self): self.assertEqual(self.model.current_val_trajectory, 0) self.assertIsNone(self.model.last_val_prediction) - # Check that .xdmf is present - xdmf_path = os.path.join( - "meshes", - f"epoch_{self.model.current_epoch}", - f"graph_epoch_{self.model.current_epoch}_{self.model.current_val_trajectory}.xdmf", - ) - h5_path = os.path.join( - "meshes", - f"epoch_{self.model.current_epoch}", - f"graph_epoch_{self.model.current_epoch}_{self.model.current_val_trajectory}.h5", - ) - - self.assertTrue(os.path.exists(xdmf_path)) - self.assertTrue(os.path.exists(h5_path)) - - _ = meshio.xdmf.TimeSeriesReader(xdmf_path) - def test_validation_step_resets_trajectory(self): # Create mock batches self.dataloader = DataLoader(self.dataset, batch_size=1) @@ -333,9 +330,6 @@ def test_prediction_step(self): self.model.eval() self.model.predict_step(batch.to(device)) - # Check that prediction_trajectory is set - self.assertIsNotNone(self.model.prediction_trajectory) - # Check that last_pred_prediction is set self.assertIsNotNone(self.model.last_pred_prediction) self.assertEqual(self.model.last_pred_prediction.shape, (10, 3)) @@ -343,6 +337,24 @@ def test_prediction_step(self): # Check that last_previous_data_pred_prediction is not set self.assertIsNone(self.model.last_previous_data_pred_prediction) + # Check that prediction files are saved + traj_idx = 0 + xdmf_path = os.path.join( + "predictions", + f"graph_{traj_idx}.xdmf", + ) + h5_path = os.path.join( + "predictions", + f"graph_{traj_idx}.h5", + ) + + self.assertTrue(os.path.exists(xdmf_path)) + self.assertTrue(os.path.exists(h5_path)) + + # Delete created prediction files + os.remove(xdmf_path) + os.remove(h5_path) + def test_predict_step_w_previous_data(self): self.dataloader = DataLoader(self.dataset, batch_size=1) batch = next(iter(self.dataloader)) @@ -357,9 +369,6 @@ def test_predict_step_w_previous_data(self): self.model.eval() self.model.predict_step(batch.to(device)) - # Check that prediction_trajectory is set - self.assertIsNotNone(self.model.prediction_trajectory) - # Check that last_pred_prediction is set self.assertIsNotNone(self.model.last_pred_prediction) self.assertEqual(self.model.last_pred_prediction.shape, (10, 3)) @@ -374,42 +383,36 @@ def test_predict_step_w_previous_data(self): self.model.previous_data_start = None self.model.previous_data_end = None + # Check that prediction files are saved + traj_idx = 0 + xdmf_path = os.path.join( + "predictions", + f"graph_{traj_idx}.xdmf", + ) + h5_path = os.path.join( + "predictions", + f"graph_{traj_idx}.h5", + ) + + self.assertTrue(os.path.exists(xdmf_path)) + self.assertTrue(os.path.exists(h5_path)) + + # Delete created prediction files + os.remove(xdmf_path) + os.remove(h5_path) + def test_on_predict_epoch_end(self): - # Simulate prediction_trajectory with sample graphs - num_graphs = 3 - for i in range(num_graphs): - # Create a simple graph - pos = torch.tensor( - [[0.0 + i, 0.0], [1.0 + i, 0.0], [1.0 + i, 1.0], [0.0 + i, 1.0]], - dtype=torch.float, - ) - edge_index = torch.tensor( - [[0, 1, 2, 3], [1, 2, 3, 0]], dtype=torch.long - ) - x = torch.tensor( - [ - [i * 10 + 1, i * 10 + 1], - [i * 10 + 2, i * 10 + 2], - [i * 10 + 3, i * 10 + 3], - [i * 10 + 4, i * 10 + 4], - ], - dtype=torch.float, - ) - face = torch.tensor([[0], [1], [2]]) - graph = Data( - pos=pos, - edge_index=edge_index, - x=x, - face=face, - ) - self.model.prediction_trajectory.append(graph) + # Make a prediction step + self.dataloader = DataLoader(self.dataset, batch_size=1) + batch = next(iter(self.dataloader)) + self.model.eval() + self.model.predict_step(batch.to(device)) # Run on_validation_epoch_end self.model.on_predict_epoch_end() - # Check that prediction_trajectory is cleared + # Check that predictions are cleared self.assertEqual(self.model.current_pred_trajectory, 0) - self.assertEqual(len(self.model.prediction_trajectory), 0) self.assertIsNone(self.model.last_pred_prediction) self.assertIsNone(self.model.last_previous_data_pred_prediction) @@ -427,38 +430,23 @@ def test_on_predict_epoch_end(self): self.assertTrue(os.path.exists(xdmf_path)) self.assertTrue(os.path.exists(h5_path)) + # Delete created prediction files + os.remove(xdmf_path) + os.remove(h5_path) + def test_on_predict_epoch_end_with_traj_id(self): - # Simulate prediction_trajectory with sample graphs that includ an ID - num_graphs = 3 - for i in range(num_graphs): - # Create a simple graph - pos = torch.tensor( - [[0.0 + i, 0.0], [1.0 + i, 0.0], [1.0 + i, 1.0], [0.0 + i, 1.0]], - dtype=torch.float, - ) - edge_index = torch.tensor( - [[0, 1, 2, 3], [1, 2, 3, 0]], dtype=torch.long - ) - x = torch.tensor( - [ - [i * 10 + 1, i * 10 + 1], - [i * 10 + 2, i * 10 + 2], - [i * 10 + 3, i * 10 + 3], - [i * 10 + 4, i * 10 + 4], - ], - dtype=torch.float, - ) - face = torch.tensor([[0], [1], [2]]) - traj_id = torch.tensor([123]) - graph = Data(pos=pos, edge_index=edge_index, x=x, face=face, id=traj_id) - self.model.prediction_trajectory.append(graph) + # Make a prediction step with a batch having an id + self.dataloader = DataLoader(self.dataset, batch_size=1) + batch = next(iter(self.dataloader)) + batch.id = torch.tensor([123]) + self.model.eval() + self.model.predict_step(batch.to(device)) # Run on_validation_epoch_end self.model.on_predict_epoch_end() - # Check that prediction_trajectory is cleared + # Check that predictions are cleared self.assertEqual(self.model.current_pred_trajectory, 0) - self.assertEqual(len(self.model.prediction_trajectory), 0) self.assertIsNone(self.model.last_pred_prediction) self.assertIsNone(self.model.last_previous_data_pred_prediction) @@ -476,6 +464,10 @@ def test_on_predict_epoch_end_with_traj_id(self): self.assertTrue(os.path.exists(xdmf_path)) self.assertTrue(os.path.exists(h5_path)) + # Delete created prediction files + os.remove(xdmf_path) + os.remove(h5_path) + def test_predict_step_saves_and_resets_trajectory(self): # Create mock batches self.dataloader = DataLoader(self.dataset, batch_size=1) @@ -490,7 +482,8 @@ def test_predict_step_saves_and_resets_trajectory(self): batch.traj_index = 2 self.model.predict_step(batch) - # Check that trajectory is saved and traj index changed + # Check that all trajectories are saved and traj index changed + xdmf_path = os.path.join( "predictions", "graph_1.xdmf", @@ -501,8 +494,23 @@ def test_predict_step_saves_and_resets_trajectory(self): ) self.assertTrue(os.path.exists(xdmf_path)) self.assertTrue(os.path.exists(h5_path)) + os.remove(xdmf_path) + os.remove(h5_path) + + xdmf_path = os.path.join( + "predictions", + "graph_2.xdmf", + ) + h5_path = os.path.join( + "predictions", + "graph_2.h5", + ) + self.assertTrue(os.path.exists(xdmf_path)) + self.assertTrue(os.path.exists(h5_path)) + os.remove(xdmf_path) + os.remove(h5_path) + assert self.model.current_pred_trajectory == 2 - # traj 2 is not saved until predict epoch end def test_wandb_run_id_on_checkpoint_save_and_load(self): self.model.wandb_run_id = "saved_run_id" diff --git a/tests/graphphysics/utils/test_meshio_mesh.py b/tests/graphphysics/utils/test_meshio_mesh.py index 063c463..efc787a 100644 --- a/tests/graphphysics/utils/test_meshio_mesh.py +++ b/tests/graphphysics/utils/test_meshio_mesh.py @@ -6,7 +6,12 @@ import torch from torch_geometric.data import Data -from graphphysics.utils.meshio_mesh import convert_to_meshio_vtu, vtu_to_xdmf +from graphphysics.utils.meshio_mesh import ( + convert_to_meshio_vtu, + vtu_to_xdmf, + meshes_to_xdmf, + append_mesh_to_xdmf, +) from tests.mock import MOCK_VTU_FOLDER_PATH, MOCK_VTU_ANEURYSM_FOLDER_PATH @@ -189,5 +194,176 @@ def test_remove_vtus(self): os.remove(f"{self.filename}.xdmf") +class TestMeshesToXdmf(unittest.TestCase): + def setUp(self): + self.meshes_2d = [ + meshio.read(os.path.join(MOCK_VTU_FOLDER_PATH, f)) + for f in os.listdir(MOCK_VTU_FOLDER_PATH) + ] + self.meshes_3d = [ + meshio.read(os.path.join(MOCK_VTU_ANEURYSM_FOLDER_PATH, f)) + for f in os.listdir(MOCK_VTU_ANEURYSM_FOLDER_PATH) + ] + + self.tmp_dir = "tests/mock_vtu_tmp" + self.filename = os.path.join(self.tmp_dir, "test_xdmf_compression") + shutil.copytree(MOCK_VTU_FOLDER_PATH, self.tmp_dir) + self.tmp_files = [ + os.path.join(self.tmp_dir, f) for f in os.listdir(self.tmp_dir) + ] + + def tearDown(self): + shutil.rmtree(self.tmp_dir) + + def test_2d_meshes(self): + """Test 2D meshes compression""" + meshes_to_xdmf(self.filename, self.meshes_2d) + + self.assertTrue(os.path.exists(f"{self.filename}.h5")) + self.assertTrue(os.path.exists(f"{self.filename}.xdmf")) + + with meshio.xdmf.TimeSeriesReader(f"{self.filename}.xdmf") as reader: + points, cells = reader.read_points_cells() + self.assertEqual(len(points), len(self.meshes_2d[0].points)) + self.assertEqual(reader.num_steps, len(self.meshes_2d)) + for i in range(reader.num_steps): + time, point_data, cell_data = reader.read_data(i) + self.assertEqual(point_data.keys(), self.meshes_2d[i].point_data.keys()) + for key in point_data.keys(): + self.assertTrue( + np.array_equal( + point_data[key], self.meshes_2d[i].point_data[key] + ) + ) + + os.remove(f"{self.filename}.h5") + os.remove(f"{self.filename}.xdmf") + + def test_3d_meshes(self): + """Test 3D meshes compression""" + + meshes_to_xdmf(self.filename, self.meshes_3d) + + self.assertTrue(os.path.exists(f"{self.filename}.h5")) + self.assertTrue(os.path.exists(f"{self.filename}.xdmf")) + + with meshio.xdmf.TimeSeriesReader(f"{self.filename}.xdmf") as reader: + points, cells = reader.read_points_cells() + self.assertEqual(len(points), len(self.meshes_3d[0].points)) + self.assertEqual(reader.num_steps, len(self.meshes_3d)) + + for i in range(reader.num_steps): + time, point_data, cell_data = reader.read_data(i) + self.assertEqual(point_data.keys(), self.meshes_3d[i].point_data.keys()) + for key in point_data.keys(): + self.assertTrue( + np.array_equal( + point_data[key], self.meshes_3d[i].point_data[key] + ) + ) + + os.remove(f"{self.filename}.h5") + os.remove(f"{self.filename}.xdmf") + + +class TestAppendMeshToXdmf(unittest.TestCase): + def setUp(self): + self.meshes_2d = [ + meshio.read(os.path.join(MOCK_VTU_FOLDER_PATH, f)) + for f in os.listdir(MOCK_VTU_FOLDER_PATH) + ] + self.meshes_3d = [ + meshio.read(os.path.join(MOCK_VTU_ANEURYSM_FOLDER_PATH, f)) + for f in os.listdir(MOCK_VTU_ANEURYSM_FOLDER_PATH) + ] + + self.tmp_dir = "tests/mock_vtu_tmp" + self.filename = os.path.join(self.tmp_dir, "test_xdmf_compression") + shutil.copytree(MOCK_VTU_FOLDER_PATH, self.tmp_dir) + self.tmp_files = [ + os.path.join(self.tmp_dir, f) for f in os.listdir(self.tmp_dir) + ] + + def tearDown(self): + shutil.rmtree(self.tmp_dir) + + def test_append_2d_mesh(self,): + """Test 2D meshes compression, adding meshes one at a time""" + meshes_to_xdmf(self.filename, [self.meshes_2d[0]]) + for mesh in self.meshes_2d[1:]: + append_mesh_to_xdmf(self.filename, mesh) + + self.assertTrue(os.path.exists(f"{self.filename}.h5")) + self.assertTrue(os.path.exists(f"{self.filename}.xdmf")) + + with meshio.xdmf.TimeSeriesReader(f"{self.filename}.xdmf") as reader: + points, cells = reader.read_points_cells() + self.assertEqual(len(points), len(self.meshes_2d[0].points)) + self.assertEqual(reader.num_steps, len(self.meshes_2d)) + for i in range(reader.num_steps): + time, point_data, cell_data = reader.read_data(i) + self.assertEqual(point_data.keys(), self.meshes_2d[i].point_data.keys()) + for key in point_data.keys(): + self.assertTrue( + np.array_equal( + point_data[key], self.meshes_2d[i].point_data[key] + ) + ) + + os.remove(f"{self.filename}.h5") + os.remove(f"{self.filename}.xdmf") + + def test_append_3d_mesh(self,): + """Test 3D meshes compression, adding meshes one at a time""" + meshes_to_xdmf(self.filename, [self.meshes_3d[0]]) + for mesh in self.meshes_3d[1:]: + append_mesh_to_xdmf(self.filename, mesh) + + self.assertTrue(os.path.exists(f"{self.filename}.h5")) + self.assertTrue(os.path.exists(f"{self.filename}.xdmf")) + + with meshio.xdmf.TimeSeriesReader(f"{self.filename}.xdmf") as reader: + points, cells = reader.read_points_cells() + self.assertEqual(len(points), len(self.meshes_3d[0].points)) + self.assertEqual(reader.num_steps, len(self.meshes_3d)) + for i in range(reader.num_steps): + time, point_data, cell_data = reader.read_data(i) + self.assertEqual(point_data.keys(), self.meshes_3d[i].point_data.keys()) + for key in point_data.keys(): + self.assertTrue( + np.array_equal( + point_data[key], self.meshes_3d[i].point_data[key] + ) + ) + + os.remove(f"{self.filename}.h5") + os.remove(f"{self.filename}.xdmf") + + def test_append_mesh_with_compression(self,): + meshes_to_xdmf(self.filename, [self.meshes_2d[0]]) + for mesh in self.meshes_2d[1:]: + append_mesh_to_xdmf(self.filename, mesh, compress=True) + + self.assertTrue(os.path.exists(f"{self.filename}.h5")) + self.assertTrue(os.path.exists(f"{self.filename}.xdmf")) + + with meshio.xdmf.TimeSeriesReader(f"{self.filename}.xdmf") as reader: + points, cells = reader.read_points_cells() + self.assertEqual(len(points), len(self.meshes_2d[0].points)) + self.assertEqual(reader.num_steps, len(self.meshes_2d)) + for i in range(reader.num_steps): + time, point_data, cell_data = reader.read_data(i) + self.assertEqual(point_data.keys(), self.meshes_2d[i].point_data.keys()) + for key in point_data.keys(): + self.assertTrue( + np.array_equal( + point_data[key], self.meshes_2d[i].point_data[key] + ) + ) + + os.remove(f"{self.filename}.h5") + os.remove(f"{self.filename}.xdmf") + + if __name__ == "__main__": unittest.main() From 9a184459aa04dc0bb10ba486dd452e3e2d90a3f2 Mon Sep 17 00:00:00 2001 From: Vincent Lannelongue Date: Tue, 24 Mar 2026 11:34:11 +0100 Subject: [PATCH 3/6] Add lxml requirement to CI/CD --- .github/workflows/gp.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/gp.yaml b/.github/workflows/gp.yaml index 3fda919..ece574f 100644 --- a/.github/workflows/gp.yaml +++ b/.github/workflows/gp.yaml @@ -39,6 +39,7 @@ jobs: pip install "wandb[media]" pip install panel pip install einops + pip install lxml - name: Linting code run: | make lint From 541d92fb7dde56dc9dabd7fa94ef61946d4be45b Mon Sep 17 00:00:00 2001 From: Vincent Lannelongue Date: Tue, 24 Mar 2026 11:43:16 +0100 Subject: [PATCH 4/6] Add compression parameter in lightning module --- graphphysics/training/lightning_module.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/graphphysics/training/lightning_module.py b/graphphysics/training/lightning_module.py index 50001aa..9475fb5 100644 --- a/graphphysics/training/lightning_module.py +++ b/graphphysics/training/lightning_module.py @@ -115,6 +115,11 @@ def __init__( self.current_pred_trajectory = 0 self.last_pred_prediction = None self.last_previous_data_pred_prediction = None + self.compress_predictions = ( + parameters["compression"] + if hasattr(self.parameters, "compression") + else False + ) training_params: Dict = parameters.get("training", {}) self.use_spatial_mtp: bool = training_params.get("use_spatial_mtp", False) @@ -146,9 +151,6 @@ def __init__( if self.use_spatial_mtp: self._setup_spatial_mtp(processor, device) - # TODO: Decide on whether or not to keep this - # self.compress_predictions = parameters["compression"] - def forward(self, graph: Batch): return self.model(graph) @@ -350,7 +352,12 @@ def _save_batch_to_xdmf( ): meshes_to_xdmf(filename=archive_path, meshes=[mesh], timestep=timestep) else: - append_mesh_to_xdmf(filename=archive_path, mesh=mesh, timestep=timestep) + append_mesh_to_xdmf( + filename=archive_path, + mesh=mesh, + timestep=timestep, + compress=self.compress_predictions, + ) def _reset_validation_trajectory(self): self.current_val_trajectory += 1 From 4617207bb0c24d592f81617387323ee3f2bb4a4c Mon Sep 17 00:00:00 2001 From: Vincent Lannelongue Date: Tue, 24 Mar 2026 11:43:38 +0100 Subject: [PATCH 5/6] Rename compression parameter --- graphphysics/training/lightning_module.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/graphphysics/training/lightning_module.py b/graphphysics/training/lightning_module.py index 9475fb5..fc64d8c 100644 --- a/graphphysics/training/lightning_module.py +++ b/graphphysics/training/lightning_module.py @@ -116,8 +116,8 @@ def __init__( self.last_pred_prediction = None self.last_previous_data_pred_prediction = None self.compress_predictions = ( - parameters["compression"] - if hasattr(self.parameters, "compression") + parameters["compress_predictions"] + if hasattr(self.parameters, "compress_predictions") else False ) From 22ae1b5d2c0c66c0ec4b0781d3366e31ac3ce7bd Mon Sep 17 00:00:00 2001 From: Vincent Lannelongue Date: Fri, 27 Mar 2026 12:00:45 +0100 Subject: [PATCH 6/6] Update readme --- README.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/README.md b/README.md index d50a865..97d2b71 100644 --- a/README.md +++ b/README.md @@ -283,6 +283,12 @@ Second, in the case of dealing with multiple meshes, you can add extra edges bas See the [description](https://arxiv.org/abs/2010.03409) regarding world edges. +Third, you can decide on adding extra compression when saving your results during validation or prediction to XDMF/HDF5 archive. This will slow down the results archiving but lower the h5 file size: + +```json +"compress_predictions": true +``` + Finally, in the case where: - you need to build the node type - you need to build extra features that were not in your dataset