Add VTKHDF output support for all structured mesh types (#3620)

openmc/mesh.py

@@ -664,12 +664,16 @@ def write_data_to_vtk(self,
                           datasets: dict | None = None,
                           volume_normalization: bool = True,
                           curvilinear: bool = False):
-        """Creates a VTK object of the mesh
+        """Creates a VTK object of the mesh and writes it to a file.
+
+        Supported formats are legacy ASCII ``.vtk`` (requires the ``vtk``
+        package) and ``.vtkhdf`` (requires only ``h5py``).
 
         Parameters
         ----------
-        filename : str
-            Name of the VTK file to write.
+        filename : str or PathLike
+            Name of the VTK file to write. Use a ``.vtkhdf`` extension for
+            the VTKHDF format or ``.vtk`` for the legacy ASCII format.
         datasets : dict
             Dictionary whose keys are the data labels and values are the data
             sets. 1D datasets are expected to be extracted directly from
@@ -680,7 +684,7 @@ def write_data_to_vtk(self,
             :class:`~openmc.MeshFilter`'s.
         volume_normalization : bool, optional
             Whether or not to normalize the data by the volume of the mesh
-            elements.
+            elements. Defaults to True.
         curvilinear : bool
             Whether or not to write curvilinear elements. Only applies to
             ``SphericalMesh`` and ``CylindricalMesh``.
@@ -692,8 +696,9 @@ def write_data_to_vtk(self,
 
         Returns
         -------
-        vtk.StructuredGrid or vtk.UnstructuredGrid
-            a VTK grid object representing the mesh
+        vtk.StructuredGrid or vtk.UnstructuredGrid or None
+            A VTK grid object representing the mesh for the legacy ASCII
+            format, or None for the VTKHDF format.
 
         Examples
         --------
@@ -712,6 +717,11 @@ def write_data_to_vtk(self,
            >>> heating = tally.get_reshaped_data(expand_dims=True)
            >>> mesh.write_data_to_vtk({'heating': heating})
         """
+        if Path(filename).suffix == ".vtkhdf":
+            self._write_vtk_hdf5(filename, datasets, volume_normalization)
+            return None
+
+        # vtk is an optional dependency only needed for the legacy ASCII path
         import vtk
         from vtk.util import numpy_support as nps
 
@@ -929,12 +939,21 @@ def _check_vtk_dataset(self, label: str, dataset: np.ndarray):
         if dataset.ndim == 1:
             return
 
-        if dataset.shape != self.dimension:
-            raise ValueError(
-                f'Cannot apply multidimensional dataset "{label}" with '
-                f"shape {dataset.shape} to mesh {self.id} "
-                f"with dimensions {self.dimension}"
-            )
+        if dataset.shape == self.dimension:
+            return
+
+        # allow datasets in lower dimension when the missing dimensions are
+        # singleton (e.g., 2D data on a 3D mesh with one depth layer)
+        if dataset.ndim < len(self.dimension):
+            padded_shape = tuple(dataset.shape) + (1,) * (len(self.dimension) - dataset.ndim)
+            if padded_shape == tuple(self.dimension):
+                return
+
+        raise ValueError(
+            f'Cannot apply multidimensional dataset "{label}" with '
+            f"shape {dataset.shape} to mesh {self.id} "
+            f"with dimensions {self.dimension}"
+        )
 
     @classmethod
     def from_domain(
@@ -1543,6 +1562,61 @@ def get_indices_at_coords(self, coords: Sequence[float]) -> tuple:
         indices = np.floor((coords_array - lower_left) / spacing).astype(int)
         return tuple(int(i) for i in indices[:ndim])
 
+    def _write_vtk_hdf5(self, filename: PathLike, datasets: dict | None, volume_normalization: bool) -> None:
+        """Write RegularMesh as VTKHDF StructuredGrid format."""
+        dims = self.dimension
+        ndim = len(dims)
+
+        # Vertex dimensions (cells + 1) – store only ndim entries so that
+        # 1-D and 2-D meshes carry the right number of dimensions.
+        vertex_dims = [d + 1 for d in dims]
+
+        # Build explicit point coordinates.  Pad coordinate arrays to 3-D so
+        # that every point has an (x, y, z) triple; extra coordinates are 0.
+        coords_1d = []
+        for i in range(ndim):
+            c = np.linspace(self.lower_left[i], self.upper_right[i], dims[i] + 1)
+            coords_1d.append(c)
+        while len(coords_1d) < 3:
+            coords_1d.append(np.array([0.0]))
+
+        # np.meshgrid with indexing='ij' → axis 0 = x, axis 1 = y, axis 2 = z
+        vertices = np.stack(
+            np.meshgrid(*coords_1d, indexing='ij'), axis=-1
+        )
+        # Swap first and last spatial axes then flatten, matching the
+        # approach used by RectilinearMesh/CylindricalMesh/SphericalMesh.
+        points = np.swapaxes(vertices, 0, 2).reshape(-1, 3).astype(np.float64)
+
+        with h5py.File(filename, "w") as f:
+            root = f.create_group("VTKHDF")
+            root.attrs["Version"] = (2, 1)
+            _type = "StructuredGrid".encode("ascii")
+            root.attrs.create(
+                "Type",
+                _type,
+                dtype=h5py.string_dtype("ascii", len(_type)),
+            )
+            root.create_dataset("Dimensions", data=vertex_dims, dtype="i8")
+            root.create_dataset("Points", data=points, dtype="f8")
+
+            cell_data_group = root.create_group("CellData")
+
+            if not datasets:
+                return
+
+            for name, data in datasets.items():
+                data = self._reshape_vtk_dataset(data)
+                self._check_vtk_dataset(name, data)
+
+                if volume_normalization:
+                    data = data / self.volumes
+
+                flat_data = data.T.ravel() if data.ndim > 1 else data.ravel()
+                cell_data_group.create_dataset(
+                    name, data=flat_data, dtype="float64", chunks=True
+                )
+
 
 def Mesh(*args, **kwargs):
     warnings.warn("Mesh has been renamed RegularMesh. Future versions of "
@@ -1796,6 +1870,50 @@ def get_indices_at_coords(self, coords: Sequence[float]) -> tuple[int, int, int]
 
         return tuple(indices)
 
+    def _write_vtk_hdf5(self, filename: PathLike, datasets: dict | None, volume_normalization: bool) -> None:
+        """Write RectilinearMesh as VTKHDF StructuredGrid format.
+
+        Note: vtkRectilinearGrid is not part of the VTKHDF spec yet, so
+        StructuredGrid with explicit point coordinates is used instead.
+        """
+        nx, ny, nz = self.dimension
+        vertex_dims = [nx + 1, ny + 1, nz + 1]
+
+        vertices = np.stack(np.meshgrid(
+            self.x_grid, self.y_grid, self.z_grid, indexing='ij'
+        ), axis=-1)
+
+        with h5py.File(filename, "w") as f:
+            root = f.create_group("VTKHDF")
+            root.attrs["Version"] = (2, 1)
+            _type = "StructuredGrid".encode("ascii")
+            root.attrs.create(
+                "Type",
+                _type,
+                dtype=h5py.string_dtype("ascii", len(_type)),
+            )
+            root.create_dataset("Dimensions", data=vertex_dims, dtype="i8")
+
+            points = np.swapaxes(vertices, 0, 2).reshape(-1, 3)
+            root.create_dataset("Points", data=points.astype(np.float64), dtype="f8")
+
+            cell_data_group = root.create_group("CellData")
+
+            if not datasets:
+                return
+
+            for name, data in datasets.items():
+                data = self._reshape_vtk_dataset(data)
+                self._check_vtk_dataset(name, data)
+
+                if volume_normalization:
+                    data = data / self.volumes
+
+                flat_data = data.T.ravel() if data.ndim > 1 else data.ravel()
+                cell_data_group.create_dataset(
+                    name, data=flat_data, dtype="float64", chunks=True
+                )
+
     @classmethod
     def from_bounding_box(
         cls,
@@ -2274,6 +2392,48 @@ def _convert_to_cartesian(arr, origin: Sequence[float]):
         arr[..., 2] += origin[2]
         return arr
 
+    def _write_vtk_hdf5(self, filename: PathLike, datasets: dict | None, volume_normalization: bool) -> None:
+        """Write CylindricalMesh as VTKHDF StructuredGrid format."""
+        nr, nphi, nz = self.dimension
+        vertex_dims = [nr + 1, nphi + 1, nz + 1]
+
+        R, Phi, Z = np.meshgrid(self.r_grid, self.phi_grid, self.z_grid, indexing='ij')
+        X = R * np.cos(Phi) + self.origin[0]
+        Y = R * np.sin(Phi) + self.origin[1]
+        Z = Z + self.origin[2]
+        vertices = np.stack([X, Y, Z], axis=-1)
+
+        with h5py.File(filename, "w") as f:
+            root = f.create_group("VTKHDF")
+            root.attrs["Version"] = (2, 1)
+            _type = "StructuredGrid".encode("ascii")
+            root.attrs.create(
+                "Type",
+                _type,
+                dtype=h5py.string_dtype("ascii", len(_type)),
+            )
+            root.create_dataset("Dimensions", data=vertex_dims, dtype="i8")
+
+            points = np.swapaxes(vertices, 0, 2).reshape(-1, 3)
+            root.create_dataset("Points", data=points.astype(np.float64), dtype="f8")
+
+            cell_data_group = root.create_group("CellData")
+
+            if not datasets:
+                return
+
+            for name, data in datasets.items():
+                data = self._reshape_vtk_dataset(data)
+                self._check_vtk_dataset(name, data)
+
+                if volume_normalization:
+                    data = data / self.volumes
+
+                flat_data = data.T.ravel() if data.ndim > 1 else data.ravel()
+                cell_data_group.create_dataset(
+                    name, data=flat_data, dtype="float64", chunks=True
+                )
+
 
 class SphericalMesh(StructuredMesh):
     """A 3D spherical mesh
@@ -2649,6 +2809,50 @@ def get_indices_at_coords(self, coords: Sequence[float]) -> tuple:
             "get_indices_at_coords is not yet implemented for SphericalMesh"
         )
 
+    def _write_vtk_hdf5(self, filename: PathLike, datasets: dict | None, volume_normalization: bool) -> None:
+        """Write SphericalMesh as VTKHDF StructuredGrid format."""
+        nr, ntheta, nphi = self.dimension
+        vertex_dims = [nr + 1, ntheta + 1, nphi + 1]
+
+        R, Theta, Phi = np.meshgrid(
+            self.r_grid, self.theta_grid, self.phi_grid, indexing='ij'
+        )
+        X = R * np.sin(Theta) * np.cos(Phi) + self.origin[0]
+        Y = R * np.sin(Theta) * np.sin(Phi) + self.origin[1]
+        Z = R * np.cos(Theta) + self.origin[2]
+        vertices = np.stack([X, Y, Z], axis=-1)
+
+        with h5py.File(filename, "w") as f:
+            root = f.create_group("VTKHDF")
+            root.attrs["Version"] = (2, 1)
+            _type = "StructuredGrid".encode("ascii")
+            root.attrs.create(
+                "Type",
+                _type,
+                dtype=h5py.string_dtype("ascii", len(_type)),
+            )
+            root.create_dataset("Dimensions", data=vertex_dims, dtype="i8")
+
+            points = np.swapaxes(vertices, 0, 2).reshape(-1, 3)
+            root.create_dataset("Points", data=points.astype(np.float64), dtype="f8")
+
+            cell_data_group = root.create_group("CellData")
+
+            if not datasets:
+                return
+
+            for name, data in datasets.items():
+                data = self._reshape_vtk_dataset(data)
+                self._check_vtk_dataset(name, data)
+
+                if volume_normalization:
+                    data = data / self.volumes
+
+                flat_data = data.T.ravel() if data.ndim > 1 else data.ravel()
+                cell_data_group.create_dataset(
+                    name, data=flat_data, dtype="float64", chunks=True
+                )
+
 
 def require_statepoint_data(func):
     @wraps(func)
@@ -3105,6 +3309,10 @@ def _write_data_to_vtk_hdf5_format(
         datasets: dict | None = None,
         volume_normalization: bool = True,
     ):
+        """Write UnstructuredMesh as VTK-HDF5 UnstructuredGrid format.
+
+        Supports linear tetrahedra and linear hexahedra elements.
+        """
         def append_dataset(dset, array):
             """Convenience function to append data to an HDF5 dataset"""
             origLen = dset.shape[0]