Merge 66ee7c7578 into d3bc1669d3

openmc/mesh.py

@@ -661,14 +661,18 @@ class StructuredMesh(MeshBase):
     def write_data_to_vtk(self,
                           filename: PathLike,
                           datasets: dict | None = None,
-                          volume_normalization: bool = True,
+                          volume_normalization: bool | None = None,
                           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
@@ -677,9 +681,9 @@ class StructuredMesh(MeshBase):
             with structured indexing in "C" ordering. See the "expand_dims" flag
             of :meth:`~openmc.Tally.get_reshaped_data` on reshaping tally data when using
             :class:`~openmc.MeshFilter`'s.
-        volume_normalization : bool, optional
+        volume_normalization : bool or None, 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``.
@@ -691,8 +695,9 @@ class StructuredMesh(MeshBase):
 
         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
         --------
@@ -711,6 +716,14 @@ class StructuredMesh(MeshBase):
            >>> heating = tally.get_reshaped_data(expand_dims=True)
            >>> mesh.write_data_to_vtk({'heating': heating})
         """
+        if volume_normalization is None:
+            volume_normalization = True
+
+        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
 
@@ -918,23 +931,29 @@ class StructuredMesh(MeshBase):
         """
         cv.check_type('data label', label, str)
 
-        if dataset.size != self.n_elements:
-            raise ValueError(
-                f"The size of the dataset '{label}' ({dataset.size}) should be"
-                f" equal to the number of mesh cells ({self.n_elements})"
-            )
+        # If this is a multidimensional dataset, validate its shape first so
+        # that the user receives the more specific error message.
+        if dataset.ndim > 1:
+            if dataset.shape == self.dimension:
+                return
 
-        # accept a flat array as-is, assuming it is in the correct order
-        if dataset.ndim == 1:
-            return
+            if dataset.ndim < len(self.dimension):
+                padded_shape = tuple(dataset.shape) + (1,) * (len(self.dimension) - dataset.ndim)
+                if padded_shape == tuple(self.dimension):
+                    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.size != self.n_elements:
+            raise ValueError(
+                f"The size of the dataset '{label}' ({dataset.size}) should be"
+                f" equal to the number of mesh cells ({self.n_elements})"
+            )
+
     @classmethod
     def from_domain(
         cls,
@@ -1133,7 +1152,7 @@ class RegularMesh(StructuredMesh):
             if self._lower_left is not None and self._dimension is not None:
                 us = self._upper_right
                 ls = self._lower_left
-                dims =  self._dimension
+                dims = self._dimension
                 return [(u - l) / d for u, l, d in zip(us, ls, dims)]
 
     @width.setter
@@ -1542,6 +1561,61 @@ class RegularMesh(StructuredMesh):
         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 "
@@ -1795,6 +1869,50 @@ class RectilinearMesh(StructuredMesh):
 
         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,
@@ -2273,6 +2391,48 @@ class CylindricalMesh(StructuredMesh):
         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
@@ -2628,7 +2788,6 @@ class SphericalMesh(StructuredMesh):
         """
         return super().centroids
 
-
     @staticmethod
     def _convert_to_cartesian(arr, origin: Sequence[float]):
         """Converts an array with r, theta, phi values in the last dimension (shape (..., 3))
@@ -2720,6 +2879,50 @@ class SphericalMesh(StructuredMesh):
 
         return (r_index, theta_index, phi_index)
 
+    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)
@@ -3179,6 +3382,18 @@ class UnstructuredMesh(MeshBase):
         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]
+            dset.resize(origLen + array.shape[0], axis=0)
+            dset[origLen:] = array
+
+        datasets = {} if datasets is None else datasets
+
         # This writer supports linear tetrahedra and linear hexahedra elements
         conn_list = []           # flattened connectivity ids
         cell_sizes = []          # number of points per cell
@@ -3209,12 +3424,15 @@ class UnstructuredMesh(MeshBase):
             )
 
         connectivity = np.asarray(conn_list, dtype=np.int64)
+        n_cells = len(vtk_types)
 
-        # Offsets must be length (numCells + 1) with a leading 0 and
-        # cumulative end-indices thereafter, per VTK's layout
-        cell_sizes_arr = np.asarray(cell_sizes, dtype=np.int64)
-        offsets = np.zeros(cell_sizes_arr.size + 1, dtype=np.int64)
-        np.cumsum(cell_sizes_arr, out=offsets[1:])
+        # Offsets are the index of the first point of each cell in the
+        # flattened connectivity array.
+        if cell_sizes:
+            # VTKHDF requires n_cells + 1 offset entries per partition
+            offsets = np.concatenate(([0], np.cumsum(cell_sizes, dtype=np.int64)))
+        else:
+            offsets = np.array([0], dtype=np.int64)
 
         for name, data in datasets.items():
             if data.shape != self.dimension:
@@ -3225,7 +3443,6 @@ class UnstructuredMesh(MeshBase):
                 )
 
         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
@@ -3236,22 +3453,37 @@ class UnstructuredMesh(MeshBase):
                 dtype=h5py.string_dtype("ascii", len(ascii_type)),
             )
 
-            # Create HDF5 file structure compliant with VTKHDF UnstructuredGrid
-            n_points = int(len(self.vertices))
-            n_cells = int(len(cell_sizes))
-            n_conn_ids = int(len(connectivity))
+            # create hdf5 file structure
+            root.create_dataset("NumberOfPoints", (0,),
+                                maxshape=(None,), dtype="i8")
+            root.create_dataset("Types", (0,), maxshape=(None,), dtype="uint8")
+            root.create_dataset(
+                "Points", (0, 3), maxshape=(None, 3), dtype="f")
+            root.create_dataset(
+                "NumberOfConnectivityIds", (0,), maxshape=(None,), dtype="i8"
+            )
+            root.create_dataset("NumberOfCells", (0,),
+                                maxshape=(None,), dtype="i8")
+            root.create_dataset("Offsets", (0,), maxshape=(None,), dtype="i8")
+            root.create_dataset("Connectivity", (0,),
+                                maxshape=(None,), dtype="i8")
 
-            root.create_dataset("NumberOfPoints", data=(n_points,), dtype="i8")
-            root.create_dataset("NumberOfCells", data=(n_cells,), dtype="i8")
-            root.create_dataset("NumberOfConnectivityIds", data=(n_conn_ids,), dtype="i8")
-            root.create_dataset("Points", data=self.vertices.astype(np.float64, copy=False), dtype="f8")
-            root.create_dataset("Types", data=np.asarray(vtk_types, dtype=np.uint8), dtype="uint8")
-            root.create_dataset("Offsets", data=offsets.astype("i8"), dtype="i8")
-            root.create_dataset("Connectivity", data=connectivity.astype("i8"), dtype="i8")
+            append_dataset(root["NumberOfPoints"],
+                           np.array([len(self.vertices)]))
+            append_dataset(root["Points"], self.vertices)
+            append_dataset(
+                root["NumberOfConnectivityIds"],
+                np.array([len(connectivity)]),
+            )
+            append_dataset(root["Connectivity"], connectivity)
+            append_dataset(root["NumberOfCells"], np.array([n_cells]))
+            append_dataset(root["Offsets"], offsets)
+            append_dataset(root["Types"], np.asarray(vtk_types, dtype="uint8"))
 
             cell_data_group = root.create_group("CellData")
 
             for name, data in datasets.items():
+                data = np.asarray(data, dtype="float64")
                 if volume_normalization:
                     data /= self.volumes
                 cell_data_group.create_dataset(

tests/unit_tests/test_mesh.py

@@ -1032,6 +1032,417 @@ def test_regular_mesh_get_indices_at_coords():
     assert len(result_1d) == 1
     assert result_1d == (5,)
 
+def test_write_vtkhdf_regular_mesh(run_in_tmpdir):
+    """Test writing a regular mesh to VTKHDF format."""
+    mesh = openmc.RegularMesh()
+    mesh.lower_left = [-5., -5., -5.]
+    mesh.upper_right = [5., 5., 5.]
+    mesh.dimension = [2, 3, 4]
+
+    # Sample some random data and write to VTKHDF
+    rng = np.random.default_rng(42)
+    ref_data = rng.random(mesh.dimension)
+    filename = "test_regular_mesh.vtkhdf"
+    mesh.write_data_to_vtk(datasets={"data": ref_data}, filename=filename)
+
+    assert Path(filename).exists()
+
+    # Verify VTKHDF file structure
+    with h5py.File(filename, "r") as f:
+        assert "VTKHDF" in f
+        root = f["VTKHDF"]
+        assert root.attrs["Type"] == b"StructuredGrid"
+        assert tuple(root.attrs["Version"]) == (2, 1)
+        assert "Dimensions" in root
+        assert "Points" in root
+        assert "CellData" in root
+        assert "data" in root["CellData"]
+
+        # Check dimensions
+        dims = root["Dimensions"][()]
+        assert tuple(dims) == (3, 4, 5)  # dimension + 1 for each
+
+        # Check data shape
+        cell_data = root["CellData"]["data"][()]
+        assert cell_data.shape == (ref_data.size,)
+
+
+def test_write_vtkhdf_rectilinear_mesh(run_in_tmpdir):
+    """Test writing a rectilinear mesh to VTKHDF format."""
+    mesh = openmc.RectilinearMesh()
+    mesh.x_grid = np.array([0., 1., 3., 6.])
+    mesh.y_grid = np.array([-5., 0., 5.])
+    mesh.z_grid = np.array([-10., -5., 0., 5., 10.])
+
+    # Sample some random data and write to VTKHDF
+    rng = np.random.default_rng(42)
+    ref_data = rng.random(mesh.dimension)
+    filename = "test_rectilinear_mesh.vtkhdf"
+    mesh.write_data_to_vtk(datasets={"flux": ref_data}, filename=filename)
+
+    assert Path(filename).exists()
+
+    # Verify VTKHDF file structure
+    with h5py.File(filename, "r") as f:
+        assert "VTKHDF" in f
+        root = f["VTKHDF"]
+        assert "CellData" in root
+        assert "flux" in root["CellData"]
+
+        # Check data was written
+        cell_data = root["CellData"]["flux"][()]
+        assert cell_data.size == ref_data.size
+
+
+def test_write_vtkhdf_cylindrical_mesh(run_in_tmpdir):
+    """Test writing a cylindrical mesh to VTKHDF format."""
+    mesh = openmc.CylindricalMesh(
+        r_grid=[0., 1., 2., 3.],
+        phi_grid=[0., pi/2, pi, 3*pi/2, 2*pi],
+        z_grid=[-5., 0., 5.],
+        origin=[0., 0., 0.]
+    )
+
+    # Sample some random data and write to VTKHDF
+    rng = np.random.default_rng(42)
+    ref_data = rng.random(mesh.dimension)
+    filename = "test_cylindrical_mesh.vtkhdf"
+    mesh.write_data_to_vtk(datasets={"power": ref_data}, filename=filename)
+
+    assert Path(filename).exists()
+
+    # Verify VTKHDF file structure
+    with h5py.File(filename, "r") as f:
+        assert "VTKHDF" in f
+        root = f["VTKHDF"]
+        assert "CellData" in root
+        assert "power" in root["CellData"]
+
+        # Check dimensions (vertices)
+        dims = root["Dimensions"][()]
+        expected_dims = np.array([3+1, 4+1, 2+1])  # r, phi, z
+        np.testing.assert_array_equal(dims, expected_dims)
+
+
+def test_write_vtkhdf_spherical_mesh(run_in_tmpdir):
+    """Test writing a spherical mesh to VTKHDF format."""
+    mesh = openmc.SphericalMesh(
+        r_grid=[0., 1., 2., 3.],
+        theta_grid=[0., pi/4, pi/2, 3*pi/4, pi],
+        phi_grid=[0., pi/2, pi, 3*pi/2, 2*pi],
+        origin=[0., 0., 0.]
+    )
+
+    # Sample some random data and write to VTKHDF
+    rng = np.random.default_rng(42)
+    ref_data = rng.random(mesh.dimension)
+    filename = "test_spherical_mesh.vtkhdf"
+    mesh.write_data_to_vtk(datasets={"density": ref_data}, filename=filename)
+
+    assert Path(filename).exists()
+
+    # Verify VTKHDF file structure
+    with h5py.File(filename, "r") as f:
+        assert "VTKHDF" in f
+        root = f["VTKHDF"]
+        assert "Points" in root
+        assert "CellData" in root
+        assert "density" in root["CellData"]
+
+
+@pytest.mark.parametrize(
+    "mesh",
+    [
+        openmc.RegularMesh.from_domain(
+            openmc.BoundingBox([0.0, 0.0, 0.0], [2.0, 2.0, 2.0]),
+            dimension=[2, 2, 2],
+        ),
+        openmc.RectilinearMesh(),
+        openmc.CylindricalMesh(
+            r_grid=[0.0, 1.0, 2.0],
+            z_grid=[0.0, 1.0, 2.0],
+            phi_grid=[0.0, np.pi / 2, np.pi],
+        ),
+        openmc.SphericalMesh(
+            r_grid=[0.0, 1.0, 2.0],
+            theta_grid=[0.0, np.pi / 2, np.pi],
+            phi_grid=[0.0, np.pi / 2, np.pi],
+        ),
+    ],
+    ids=["regular", "rectilinear", "cylindrical", "spherical"],
+)
+def test_write_vtkhdf_structuredgrid_points_order_and_type(run_in_tmpdir, mesh):
+    """Test VTKHDF _write_vtk_hdf5 point ordering and Type attribute."""
+    if isinstance(mesh, openmc.RectilinearMesh):
+        mesh.x_grid = [0.0, 1.0, 2.0]
+        mesh.y_grid = [0.0, 1.0, 2.0]
+        mesh.z_grid = [0.0, 1.0, 2.0]
+    data = np.arange(mesh.n_elements).reshape(mesh.dimension)
+    mesh.write_data_to_vtk(
+        datasets={"values": data},
+        filename="test_structured.vtkhdf",
+        volume_normalization=False,
+    )
+
+    with h5py.File("test_structured.vtkhdf", "r") as f:
+        root = f["VTKHDF"]
+
+        points = root["Points"][()]
+        expected_points = np.swapaxes(mesh.vertices, 0, 2).reshape(-1, 3)
+        np.testing.assert_allclose(points, expected_points)
+
+        type_attr = root.attrs["Type"]
+        if isinstance(type_attr, bytes):
+            type_attr = type_attr.decode("ascii")
+        assert type_attr == "StructuredGrid"
+
+
+def test_write_vtkhdf_volume_normalization(run_in_tmpdir):
+    """Test volume normalization in VTKHDF format."""
+    mesh = openmc.RegularMesh()
+    mesh.lower_left = [0., 0., 0.]
+    mesh.upper_right = [10., 10., 10.]
+    mesh.dimension = [2, 2, 2]
+
+    # Use non-uniform data so that .T.ravel() and .ravel() differ,
+    # catching any ordering mistakes in the implementation.
+    ref_data = np.arange(8, dtype=float).reshape(mesh.dimension)
+    filename_with_norm = "test_normalized.vtkhdf"
+    filename_without_norm = "test_unnormalized.vtkhdf"
+
+    # Write with normalization
+    mesh.write_data_to_vtk(
+        datasets={"flux": ref_data},
+        filename=filename_with_norm,
+        volume_normalization=True
+    )
+
+    # Write without normalization
+    mesh.write_data_to_vtk(
+        datasets={"flux": ref_data},
+        filename=filename_without_norm,
+        volume_normalization=False
+    )
+
+    # Read both files and compare
+    with h5py.File(filename_with_norm, "r") as f:
+        normalized_data = f["VTKHDF"]["CellData"]["flux"][()]
+
+    with h5py.File(filename_without_norm, "r") as f:
+        unnormalized_data = f["VTKHDF"]["CellData"]["flux"][()]
+
+    # Volume for each cell is 5 x 5 x 5 = 125
+    cell_volume = 125.0
+    expected_normalized = ref_data.T.ravel() / cell_volume
+    np.testing.assert_allclose(normalized_data, expected_normalized)
+    np.testing.assert_allclose(unnormalized_data, ref_data.T.ravel())
+
+
+def test_write_vtkhdf_default_volume_normalization_for_vtkhdf(run_in_tmpdir):
+    """Ensure VTKHDF default normalizes data by volume (same as legacy)."""
+    mesh = openmc.RegularMesh()
+    mesh.lower_left = [0.0, 0.0, 0.0]
+    mesh.upper_right = [10.0, 10.0, 10.0]
+    mesh.dimension = [2, 2, 2]
+
+    data = np.ones(mesh.dimension) * 100.0
+    mesh.write_data_to_vtk(
+        datasets={"flux": data},
+        filename="test_default_norm.vtkhdf",
+    )
+
+    with h5py.File("test_default_norm.vtkhdf", "r") as f:
+        saved = f["VTKHDF"]["CellData"]["flux"][()]
+
+    # Each cell volume is 5*5*5 = 125, default normalization divides by it
+    cell_volume = 125.0
+    np.testing.assert_allclose(saved, (data / cell_volume).T.ravel())
+
+
+@pytest.mark.parametrize(
+    "x_grid,y_grid,z_grid,data,expected", [
+        (
+            [0.0, 1.0, 2.0],
+            [0.0, 1.0],
+            [0.0, 1.0],
+            np.arange(2),
+            np.arange(2),
+        ),
+        (
+            [0.0, 1.0, 2.0],
+            [0.0, 1.0, 2.0],
+            [0.0, 1.0],
+            np.arange(4).reshape(2,2),
+            np.arange(4),
+        ),
+        (
+            [0.0, 1.0, 2.0],
+            [0.0, 1.0, 2.0],
+            [0.0, 1.0, 2.0],
+            np.arange(8).reshape(2,2,2),
+            np.arange(8),
+        ),
+    ],
+    ids=["1d", "2d", "3d"],
+)
+def test_write_vtkhdf_rectilinear_multi_dimensional_data(run_in_tmpdir, x_grid, y_grid, z_grid, data, expected):
+    """Test RectilinearMesh VTKHDF accepts 1D/2D/3D data with singleton padding."""
+    mesh = openmc.RectilinearMesh()
+    mesh.x_grid = x_grid
+    mesh.y_grid = y_grid
+    mesh.z_grid = z_grid
+
+    mesh.write_data_to_vtk(
+        datasets={"flux": data},
+        filename="test_rectilinear_dims.vtkhdf",
+        volume_normalization=False,
+    )
+
+    with h5py.File("test_rectilinear_dims.vtkhdf", "r") as f:
+        saved = f["VTKHDF"]["CellData"]["flux"][()]
+
+    expected_written = data.T.ravel() if data.ndim > 1 else data.ravel()
+    np.testing.assert_allclose(saved, expected_written)
+
+
+def test_write_vtkhdf_multiple_datasets(run_in_tmpdir):
+    """Test writing multiple datasets to VTKHDF format."""
+    mesh = openmc.RegularMesh()
+    mesh.lower_left = [0., 0., 0.]
+    mesh.upper_right = [1., 1., 1.]
+    mesh.dimension = [2, 2, 2]
+
+    # Create multiple datasets
+    rng = np.random.default_rng(42)
+    data1 = rng.random(mesh.dimension)
+    data2 = rng.random(mesh.dimension)
+    data3 = rng.random(mesh.dimension)
+
+    filename = "test_multiple_datasets.vtkhdf"
+    mesh.write_data_to_vtk(
+        datasets={"flux": data1, "power": data2, "heating": data3},
+        filename=filename,
+        volume_normalization=False
+    )
+
+    assert Path(filename).exists()
+
+    # Verify all datasets are present
+    with h5py.File(filename, "r") as f:
+        root = f["VTKHDF"]
+        assert "flux" in root["CellData"]
+        assert "power" in root["CellData"]
+        assert "heating" in root["CellData"]
+
+        # Verify data integrity
+        np.testing.assert_allclose(
+            root["CellData"]["flux"][()],
+            data1.T.ravel()
+        )
+        np.testing.assert_allclose(
+            root["CellData"]["power"][()],
+            data2.T.ravel()
+        )
+        np.testing.assert_allclose(
+            root["CellData"]["heating"][()],
+            data3.T.ravel()
+        )
+
+
+def test_write_vtkhdf_invalid_data_shape(run_in_tmpdir):
+    """Test that VTKHDF raises error for mismatched data shape."""
+    mesh = openmc.RegularMesh()
+    mesh.lower_left = [0., 0., 0.]
+    mesh.upper_right = [1., 1., 1.]
+    mesh.dimension = [2, 2, 2]
+
+    # Create data with wrong shape but correct total size
+    wrong_data = np.ones((4, 2, 1))
+    filename = "test_invalid_shape.vtkhdf"
+
+    with pytest.raises(ValueError, match="Cannot apply multidimensional dataset"):
+        mesh.write_data_to_vtk(
+            datasets={"bad_data": wrong_data},
+            filename=filename
+        )
+
+
+def test_write_vtkhdf_1d_mesh(run_in_tmpdir):
+    """Test writing a 1D regular mesh to VTKHDF format."""
+    mesh = openmc.RegularMesh()
+    mesh.lower_left = [0.]
+    mesh.upper_right = [10.]
+    mesh.dimension = [5]
+
+    ref_data = np.arange(5, dtype=float)
+    filename = "test_1d_mesh.vtkhdf"
+    mesh.write_data_to_vtk(datasets={"data": ref_data}, filename=filename)
+
+    assert Path(filename).exists()
+
+    with h5py.File(filename, "r") as f:
+        root = f["VTKHDF"]
+        dims = root["Dimensions"][()]
+        assert len(dims) == 1
+        saved = root["CellData"]["data"][()]
+
+    # Default volume_normalization=True; cell width = 2.0 → volume = 2.0
+    np.testing.assert_allclose(saved, ref_data / 2.0)
+
+
+def test_write_vtkhdf_2d_mesh(run_in_tmpdir):
+    """Test writing a 2D regular mesh to VTKHDF format."""
+    mesh = openmc.RegularMesh()
+    mesh.lower_left = [0., 0.]
+    mesh.upper_right = [10., 10.]
+    mesh.dimension = [5, 3]
+
+    ref_data = np.arange(15, dtype=float).reshape(mesh.dimension)
+    filename = "test_2d_mesh.vtkhdf"
+    mesh.write_data_to_vtk(datasets={"data": ref_data}, filename=filename)
+
+    assert Path(filename).exists()
+
+    with h5py.File(filename, "r") as f:
+        root = f["VTKHDF"]
+        dims = root["Dimensions"][()]
+        assert len(dims) == 2
+        saved = root["CellData"]["data"][()]
+
+    # Default volume_normalization=True; cell volume = 2.0 * (10/3)
+    cell_volume = 2.0 * (10.0 / 3.0)
+    np.testing.assert_allclose(saved, (ref_data / cell_volume).T.ravel())
+
+
+def test_write_ascii_vtk_unchanged(run_in_tmpdir):
+    """Test that ASCII .vtk format still works as before."""
+    mesh = openmc.RegularMesh()
+    mesh.lower_left = [0., 0., 0.]
+    mesh.upper_right = [1., 1., 1.]
+    mesh.dimension = [2, 2, 2]
+
+    rng = np.random.default_rng(42)
+    ref_data = rng.random(mesh.dimension)
+    filename = "test_ascii.vtk"
+
+    # This should work without requiring vtk module changes
+    vtkIOLegacy = pytest.importorskip("vtkmodules.vtkIOLegacy")
+    mesh.write_data_to_vtk(datasets={"data": ref_data},
+                           filename=filename,
+                           volume_normalization=False
+                           )
+
+    assert Path(filename).exists()
+
+    # Verify with VTK reader
+    reader = vtkIOLegacy.vtkGenericDataObjectReader()
+    reader.SetFileName(str(filename))
+    reader.Update()
+
+    # Get data from file
+    arr = reader.GetOutput().GetCellData().GetArray("data")
+    read_data = np.array([arr.GetTuple1(i) for i in range(ref_data.size)])
+    np.testing.assert_almost_equal(read_data, ref_data.T.ravel())
 
 def test_rectilinear_mesh_get_indices_at_coords():
     """Test get_indices_at_coords method for RectilinearMesh"""