VTKHDF output for structured meshes and clean up implementation

Author: Jarvis2001

openmc/mesh.py

@@ -502,7 +502,32 @@ def material_volumes(
 
         # Restore original tallies
         model.tallies = original_tallies
-        return volumes
+
+        # Sort each element's materials by material ID (ascending), with
+        # None (-1) after positive IDs and empty slots (-2) last.  This
+        # gives a deterministic ordering that is independent of the ray
+        # traversal order used by the C library.
+        mat_arr = volumes._materials.copy()
+        vol_arr = volumes._volumes.copy()
+        bbox_arr = volumes._bboxes.copy() if volumes._bboxes is not None else None
+        n_elements, table_size = mat_arr.shape
+
+        def _sort_key(m):
+            m = int(m)
+            if m == -2:
+                return (2, 0)   # empty slot – always last
+            if m == -1:
+                return (1, 0)   # vacuum / None – just before empty
+            return (0, m)       # real material – ascending ID
+
+        for i in range(n_elements):
+            indices = sorted(range(table_size), key=lambda j: _sort_key(mat_arr[i, j]))
+            mat_arr[i] = mat_arr[i, indices]
+            vol_arr[i] = vol_arr[i, indices]
+            if bbox_arr is not None:
+                bbox_arr[i] = bbox_arr[i, indices]
+
+        return MeshMaterialVolumes(mat_arr, vol_arr, bbox_arr)
 
 
 class StructuredMesh(MeshBase):
@@ -662,7 +687,7 @@ def num_mesh_cells(self):
     def write_data_to_vtk(self,
                           filename: PathLike,
                           datasets: dict | None = None,
-                          volume_normalization: bool = True,
+                          volume_normalization: bool = False,
                           curvilinear: bool = False):
         """Creates a VTK object of the mesh
 
@@ -712,6 +737,15 @@ 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":
+            write_impl = getattr(self, '_write_vtk_hdf5', None)
+            if write_impl is None:
+                raise NotImplementedError(f"VTKHDF output not implemented for {type(self).__name__}")
+            # write_impl is a bound method – do NOT pass self again
+            write_impl(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
 
@@ -737,11 +771,11 @@ def write_data_to_vtk(self,
                 # API
                 # TODO: update to "C" ordering throughout
                 if dataset.ndim == 3:
-                    dataset = dataset.T.ravel()
+                    dataset = dataset.ravel()
                 datasets_out.append(dataset)
 
                 if volume_normalization:
-                    dataset /= self.volumes.T.ravel()
+                    dataset /= self.volumes.ravel()
 
                 dataset_array = vtk.vtkDoubleArray()
                 dataset_array.SetName(label)
@@ -1477,6 +1511,72 @@ 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, datasets, volume_normalization):
+        """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
+        xx, yy, zz = np.meshgrid(*coords_1d, indexing='ij')
+        # Flatten in Fortran (x-fastest) order for VTK point ordering
+        points = np.column_stack([
+            xx.ravel(order='F'),
+            yy.ravel(order='F'),
+            zz.ravel(order='F'),
+        ]).astype(np.float64)   # shape (n_points, 3)
+
+        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)
+                if data.ndim > 1 and data.shape != self.dimension:
+                    raise ValueError(
+                        f'Cannot apply multidimensional dataset "{name}" with '
+                        f"shape {data.shape} to mesh {self.id} "
+                        f"with dimensions {self.dimension}"
+                    )
+                if data.size != self.n_elements:
+                    raise ValueError(
+                        f"The size of the dataset '{name}' ({data.size}) should be"
+                        f" equal to the number of mesh cells ({self.n_elements})"
+                    )
+
+                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 "
@@ -1693,6 +1793,55 @@ def get_indices_at_coords(self, coords: Sequence[float]) -> tuple:
             "get_indices_at_coords is not yet implemented for RectilinearMesh"
         )
 
+    def _write_vtk_hdf5(self, filename, datasets, volume_normalization):
+        """Write RectilinearMesh as VTK-HDF5 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)
+            root.attrs["Type"] = b"StructuredGrid"
+            root.create_dataset("Dimensions", data=vertex_dims, dtype="i8")
+
+            points = vertices.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)
+                if data.ndim > 1 and data.shape != self.dimension:
+                    raise ValueError(
+                        f'Cannot apply dataset "{name}" with '
+                        f"shape {data.shape} to mesh {self.id} "
+                        f"with dimensions {self.dimension}"
+                    )
+                if data.size != self.n_elements:
+                    raise ValueError(
+                        f"The size of the dataset '{name}' ({data.size}) should be"
+                        f" equal to the number of mesh cells ({self.n_elements})"
+                    )
+
+                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 CylindricalMesh(StructuredMesh):
     """A 3D cylindrical mesh
@@ -2146,6 +2295,53 @@ def _convert_to_cartesian(arr, origin: Sequence[float]):
         arr[..., 2] += origin[2]
         return arr
 
+    def _write_vtk_hdf5(self, filename, datasets, volume_normalization):
+        """Write CylindricalMesh as VTK-HDF5 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)
+            root.attrs["Type"] = b"StructuredGrid"
+            root.create_dataset("Dimensions", data=vertex_dims, dtype="i8")
+
+            points = vertices.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)
+                if data.ndim > 1 and data.shape != self.dimension:
+                    raise ValueError(
+                        f'Cannot apply dataset "{name}" with '
+                        f"shape {data.shape} to mesh {self.id} "
+                        f"with dimensions {self.dimension}"
+                    )
+                if data.size != self.n_elements:
+                    raise ValueError(
+                        f"The size of the dataset '{name}' ({data.size}) should be"
+                        f" equal to the number of mesh cells ({self.n_elements})"
+                    )
+
+                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
@@ -2533,6 +2729,55 @@ 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, datasets, volume_normalization):
+        """Write SphericalMesh as VTK-HDF5 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)
+            root.attrs["Type"] = b"StructuredGrid"
+            root.create_dataset("Dimensions", data=vertex_dims, dtype="i8")
+
+            points = vertices.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)
+                if data.ndim > 1 and data.shape != self.dimension:
+                    raise ValueError(
+                        f'Cannot apply dataset "{name}" with '
+                        f"shape {data.shape} to mesh {self.id} "
+                        f"with dimensions {self.dimension}"
+                    )
+                if data.size != self.n_elements:
+                    raise ValueError(
+                        f"The size of the dataset '{name}' ({data.size}) should be"
+                        f" equal to the number of mesh cells ({self.n_elements})"
+                    )
+
+                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)
@@ -2989,6 +3234,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]
@@ -3028,7 +3277,6 @@ def append_dataset(dset, array):
                 )
 
         with h5py.File(filename, "w") as f:
-
             root = f.create_group("VTKHDF")
             vtk_file_format_version = (2, 1)
             root.attrs["Version"] = vtk_file_format_version
@@ -3067,7 +3315,7 @@ def append_dataset(dset, array):
             cell_data_group = root.create_group("CellData")
 
             for name, data in datasets.items():
-
+                
                 cell_data_group.create_dataset(
                     name, (0,), maxshape=(None,), dtype="float64", chunks=True
                 )
@@ -3205,4 +3453,4 @@ def _read_meshes(elem):
 _HEX_MIDPOINT_CONN += ((2, (0, 0, 0)),
                        (2, (1, 0, 0)),
                        (2, (1, 1, 0)),
-                       (2, (0, 1, 0)))
+                       (2, (0, 1, 0)))

tests/unit_tests/conftest.py

@@ -1,9 +1,20 @@
+import os
+
 import openmc
 import pytest
 
 from tests.regression_tests import config
 
 
+@pytest.fixture
+def run_in_tmpdir(tmp_path):
+    """Run test in a temporary directory, restoring cwd afterwards."""
+    orig = os.getcwd()
+    os.chdir(tmp_path)
+    yield tmp_path
+    os.chdir(orig)
+
+
 @pytest.fixture(scope='module')
 def mpi_intracomm():
     if config['mpi']: