diff --git a/openmc/mesh.py b/openmc/mesh.py index 4299fc5530d..da37b090d80 100644 --- a/openmc/mesh.py +++ b/openmc/mesh.py @@ -22,6 +22,23 @@ from .utility_funcs import input_path +# VTKHDF format version used for structured mesh output +_VTKHDF_VERSION = (2, 1) + + +def _write_vtkhdf_type(group, name: str): + """Write the VTKHDF dataset type as a fixed length ASCII string + + A variable length string is read back as ``str`` rather than ``bytes``, + which does not match what VTK readers expect. + + """ + ascii_name = name.encode("ascii") + group.attrs.create( + "Type", ascii_name, dtype=h5py.string_dtype("ascii", len(ascii_name)) + ) + + class MeshMaterialVolumes(Mapping): """Results from a material volume in mesh calculation. @@ -673,12 +690,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 @@ -701,8 +722,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 -------- @@ -721,6 +743,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 @@ -928,22 +955,138 @@ def _check_vtk_dataset(self, label: str, dataset: np.ndarray): """ cv.check_type('data label', label, str) + # 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 + + 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}" + ) + 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})" ) - # accept a flat array as-is, assuming it is in the correct order - if dataset.ndim == 1: - return + def _write_vtk_hdf5(self, filename: PathLike, datasets: dict | None, + volume_normalization: bool) -> None: + """Write the mesh and data in the VTKHDF UnstructuredGrid format - 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}" + VTKHDF defines no StructuredGrid or RectilinearGrid dataset type, so + meshes with non-uniform or curvilinear elements are written as an + unstructured grid of linear hexahedra. :class:`RegularMesh` overrides + this with the more compact ImageData format. + + Parameters + ---------- + filename : str or PathLike + Name of the VTKHDF file to write. + datasets : dict or None + Dictionary whose keys are the data labels and values are the data + sets. + volume_normalization : bool + Whether or not to normalize the data by the volume of the mesh + elements. + + """ + n_i, n_j, n_k = self.dimension + vertex_shape = (n_i + 1, n_j + 1, n_k + 1) + points = np.asarray(self.vertices).reshape(-1, 3) + + # Index of the (i, j, k) corner of each element, ordered with i varying + # fastest to match the cell ordering used elsewhere for VTK output + i, j, k = np.meshgrid(np.arange(n_i), np.arange(n_j), np.arange(n_k), + indexing='ij') + corner = np.ravel_multi_index((i, j, k), vertex_shape) + corner = corner.transpose(2, 1, 0).ravel() + + # Offset from the (i, j, k) corner to each vertex of a VTK hexahedron + stride_i = (n_j + 1) * (n_k + 1) + stride_j = n_k + 1 + offsets = np.array([ + 0, + stride_i, + stride_i + stride_j, + stride_j, + 1, + stride_i + 1, + stride_i + stride_j + 1, + stride_j + 1, + ]) + connectivity = (corner[:, None] + offsets[None, :]).ravel() + + n_cells = corner.size + with h5py.File(filename, "w") as f: + root = f.create_group("VTKHDF") + root.attrs["Version"] = _VTKHDF_VERSION + _write_vtkhdf_type(root, "UnstructuredGrid") + root.create_dataset("NumberOfPoints", data=[len(points)], dtype="i8") + root.create_dataset("Points", data=points, dtype="f8") + root.create_dataset("NumberOfCells", data=[n_cells], dtype="i8") + root.create_dataset("NumberOfConnectivityIds", + data=[connectivity.size], dtype="i8") + root.create_dataset("Connectivity", data=connectivity, dtype="i8") + root.create_dataset( + "Offsets", + data=np.arange(0, connectivity.size + 1, _N_HEX_VERTICES), + dtype="i8", ) + # 12 is the VTK cell type for a linear hexahedron + root.create_dataset("Types", data=np.full(n_cells, 12), + dtype="uint8") + + cell_data_group = root.create_group("CellData") + for label, dataset in self._vtk_hdf5_datasets(datasets, + volume_normalization): + cell_data_group.create_dataset(label, data=dataset, + dtype="float64", chunks=True) + + def _vtk_hdf5_datasets(self, datasets: dict | None, + volume_normalization: bool): + """Validate and flatten datasets for VTKHDF output + + Yields + ------ + tuple of (str, numpy.ndarray) + Data label and the corresponding flat array of element values, + ordered with the first mesh index varying fastest. + + """ + if not datasets: + return + + for label, dataset in datasets.items(): + dataset = self._reshape_vtk_dataset(dataset) + self._check_vtk_dataset(label, dataset) + # restore any trailing axes of size one so that the dataset + # aligns with the volumes instead of broadcasting against them + if dataset.ndim > 1: + dataset = dataset.reshape(self.dimension) + + if volume_normalization: + # a flat dataset is already ordered with the first index + # varying fastest, so the volumes have to match that ordering + if dataset.ndim > 1: + dataset = dataset / self.volumes + else: + dataset = dataset / self.volumes.T.ravel() + + if dataset.ndim > 1: + dataset = dataset.T.ravel() + else: + dataset = dataset.ravel() + + yield label, dataset @classmethod def from_domain( @@ -1146,7 +1289,7 @@ def width(self): 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 @@ -1555,6 +1698,56 @@ 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 the mesh and data in the VTKHDF ImageData format + + A regular mesh is a uniform grid, so it is fully described by an origin + and a spacing and no point coordinates need to be stored. Meshes with + fewer than three dimensions are padded to three, as required by the + VTKHDF ImageData specification. + + Parameters + ---------- + filename : str or PathLike + Name of the VTKHDF file to write. + datasets : dict or None + Dictionary whose keys are the data labels and values are the data + sets. + volume_normalization : bool + Whether or not to normalize the data by the volume of the mesh + elements. + + """ + ndim = len(self.dimension) + shape = list(self.dimension) + [1] * (3 - ndim) + origin = list(self.lower_left) + [0.0] * (3 - ndim) + spacing = list(self.width) + [1.0] * (3 - ndim) + + with h5py.File(filename, "w") as f: + root = f.create_group("VTKHDF") + root.attrs["Version"] = _VTKHDF_VERSION + _write_vtkhdf_type(root, "ImageData") + root.attrs.create( + "WholeExtent", + np.array([0, shape[0], 0, shape[1], 0, shape[2]], dtype="i8"), + ) + root.attrs.create("Origin", np.array(origin, dtype="f8")) + root.attrs.create("Spacing", np.array(spacing, dtype="f8")) + root.attrs.create("Direction", np.eye(3).ravel()) + + cell_data_group = root.create_group("CellData") + for label, dataset in self._vtk_hdf5_datasets(datasets, + volume_normalization): + # ImageData element data is stored with the last index varying + # fastest, i.e. with shape (nz, ny, nx) + cell_data_group.create_dataset( + label, + data=dataset.reshape(shape[2], shape[1], shape[0]), + dtype="float64", + chunks=True, + ) + def Mesh(*args, **kwargs): warnings.warn("Mesh has been renamed RegularMesh. Future versions of " @@ -1810,6 +2003,7 @@ def get_indices_at_coords(self, coords: Sequence[float]) -> tuple[int, int, int] return tuple(indices) + @classmethod def from_bounding_box( cls, @@ -2291,6 +2485,7 @@ def _convert_to_cartesian(arr, origin: Sequence[float]): return arr + class SphericalMesh(StructuredMesh): """A 3D spherical mesh @@ -2647,7 +2842,6 @@ def centroids_spherical(self): """ 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)) @@ -2740,6 +2934,7 @@ def get_indices_at_coords( return (r_index, theta_index, phi_index) + def require_statepoint_data(func): @wraps(func) def wrapper(self: UnstructuredMesh, *args, **kwargs): @@ -3232,6 +3427,18 @@ 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] + 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 @@ -3262,12 +3469,15 @@ def _write_data_to_vtk_hdf5_format( ) 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: @@ -3278,7 +3488,6 @@ def _write_data_to_vtk_hdf5_format( ) 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 @@ -3289,22 +3498,37 @@ def _write_data_to_vtk_hdf5_format( 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)) - - 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") + # 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") + + 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( diff --git a/tests/unit_tests/test_mesh.py b/tests/unit_tests/test_mesh.py index 3aae6213ae9..b1e6b2b2474 100644 --- a/tests/unit_tests/test_mesh.py +++ b/tests/unit_tests/test_mesh.py @@ -1202,6 +1202,475 @@ def test_regular_mesh_get_indices_at_coords(): assert len(result_1d) == 1 assert result_1d == (5,) +def _read_vtkhdf(filename): + """Read a VTKHDF file with VTK, skipping the test if VTK is unavailable""" + vtk = pytest.importorskip("vtk") + reader = vtk.vtkHDFReader() + reader.SetFileName(str(filename)) + reader.Update() + return reader.GetOutput() + + +VTKHDF_MESHES = { + "regular": lambda: openmc.RegularMesh.from_domain( + openmc.BoundingBox([0.0, 0.0, 0.0], [2.0, 3.0, 4.0]), + dimension=[2, 3, 4], + ), + "rectilinear": lambda: openmc.RectilinearMesh.from_domain( + openmc.BoundingBox([0.0, 0.0, 0.0], [2.0, 3.0, 4.0]), + dimension=[2, 3, 4], + ), + "cylindrical": lambda: openmc.CylindricalMesh( + r_grid=[0.0, 1.0, 2.0], + phi_grid=np.linspace(0.0, 2 * pi, 4), + z_grid=[0.0, 1.0, 2.0, 3.0, 4.0], + ), + "spherical": lambda: openmc.SphericalMesh( + r_grid=[0.0, 1.0, 2.0], + theta_grid=np.linspace(0.0, pi, 4), + phi_grid=np.linspace(0.0, 2 * pi, 6), + ), +} + + +@pytest.mark.parametrize("kind", list(VTKHDF_MESHES)) +def test_write_vtkhdf_readable_by_vtk(run_in_tmpdir, kind): + """VTKHDF output must be readable by VTK, not merely well formed HDF5. + + VTKHDF has no StructuredGrid or RectilinearGrid dataset type, so a file + claiming to be one is silently rejected by every VTK reader. Assert that + the file loads and that the element data survives the round trip. + """ + mesh = VTKHDF_MESHES[kind]() + # distinct values per element so any reordering is detected + data = np.arange(float(mesh.n_elements)).reshape(mesh.dimension) + + filename = f"test_readable_{kind}.vtkhdf" + mesh.write_data_to_vtk(datasets={"values": data}, filename=filename, + volume_normalization=False) + + grid = _read_vtkhdf(filename) + assert grid.GetNumberOfCells() == mesh.n_elements + assert grid.GetNumberOfPoints() == np.prod([d + 1 for d in mesh.dimension]) + + array = grid.GetCellData().GetArray("values") + assert array is not None + read_data = np.array([array.GetTuple1(i) for i in range(mesh.n_elements)]) + np.testing.assert_allclose(read_data, data.T.ravel()) + + +@pytest.mark.parametrize("kind", list(VTKHDF_MESHES)) +def test_write_vtkhdf_element_geometry(run_in_tmpdir, kind): + """Each element read back by VTK must span the expected mesh vertices""" + mesh = VTKHDF_MESHES[kind]() + mesh.write_data_to_vtk(datasets=None, filename=f"test_geom_{kind}.vtkhdf") + + grid = _read_vtkhdf(f"test_geom_{kind}.vtkhdf") + n_i, n_j, _ = mesh.dimension + vertices = np.asarray(mesh.vertices) + + for cell_id in range(grid.GetNumberOfCells()): + i = cell_id % n_i + j = (cell_id // n_i) % n_j + k = cell_id // (n_i * n_j) + cell = grid.GetCell(cell_id) + found = np.array([cell.GetPoints().GetPoint(n) for n in range(8)]) + expected = vertices[i:i+2, j:j+2, k:k+2].reshape(-1, 3) + np.testing.assert_allclose( + np.array(sorted(map(tuple, np.round(found, 8)))), + np.array(sorted(map(tuple, np.round(expected, 8)))), + atol=1e-8, + ) + + +def test_write_vtkhdf_regular_mesh(run_in_tmpdir): + """A regular mesh is uniform, so it is written as VTKHDF ImageData.""" + mesh = openmc.RegularMesh() + mesh.lower_left = [-5., -5., -5.] + mesh.upper_right = [5., 5., 5.] + mesh.dimension = [2, 3, 4] + + 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() + + with h5py.File(filename, "r") as f: + assert "VTKHDF" in f + root = f["VTKHDF"] + assert root.attrs["Type"] == b"ImageData" + assert tuple(root.attrs["Version"]) == (2, 1) + assert "CellData" in root + assert "data" in root["CellData"] + + # ImageData is described by an extent, origin and spacing rather than + # explicit point coordinates + assert "Points" not in root + np.testing.assert_array_equal(root.attrs["WholeExtent"], + [0, 2, 0, 3, 0, 4]) + np.testing.assert_allclose(root.attrs["Origin"], [-5., -5., -5.]) + np.testing.assert_allclose(root.attrs["Spacing"], [5., 10/3, 2.5]) + + # element data is stored with the last index varying fastest + assert root["CellData"]["data"].shape == (4, 3, 2) + + +def test_write_vtkhdf_rectilinear_mesh(run_in_tmpdir): + """A rectilinear mesh is written as a VTKHDF UnstructuredGrid.""" + 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.]) + + 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() + + with h5py.File(filename, "r") as f: + assert "VTKHDF" in f + root = f["VTKHDF"] + assert root.attrs["Type"] == b"UnstructuredGrid" + assert "CellData" in root + assert "flux" in root["CellData"] + + assert root["CellData"]["flux"].size == ref_data.size + assert root["NumberOfCells"][0] == mesh.n_elements + # all elements are linear hexahedra + assert np.all(root["Types"][()] == 12) + assert root["Offsets"].size == mesh.n_elements + 1 + + +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.] + ) + + 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() + + with h5py.File(filename, "r") as f: + assert "VTKHDF" in f + root = f["VTKHDF"] + assert root.attrs["Type"] == b"UnstructuredGrid" + assert "CellData" in root + assert "power" in root["CellData"] + + # one point per mesh vertex, r x phi x z + assert root["NumberOfPoints"][0] == (3+1) * (4+1) * (2+1) + assert root["NumberOfCells"][0] == 3 * 4 * 2 + + +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.] + ) + + 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() + + with h5py.File(filename, "r") as f: + assert "VTKHDF" in f + root = f["VTKHDF"] + assert root.attrs["Type"] == b"UnstructuredGrid" + assert "Points" in root + assert "CellData" in root + assert "density" in root["CellData"] + + +@pytest.mark.parametrize( + "kind", ["rectilinear", "cylindrical", "spherical"] +) +def test_write_vtkhdf_points_match_vertices(run_in_tmpdir, kind): + """Points written for an unstructured grid must be the mesh vertices.""" + mesh = VTKHDF_MESHES[kind]() + data = np.arange(float(mesh.n_elements)).reshape(mesh.dimension) + mesh.write_data_to_vtk( + datasets={"values": data}, + filename="test_unstructured.vtkhdf", + volume_normalization=False, + ) + + with h5py.File("test_unstructured.vtkhdf", "r") as f: + root = f["VTKHDF"] + np.testing.assert_allclose( + root["Points"][()], np.asarray(mesh.vertices).reshape(-1, 3) + ) + assert root.attrs["Type"] == b"UnstructuredGrid" + + +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" + + mesh.write_data_to_vtk( + datasets={"flux": ref_data}, + filename=filename_with_norm, + volume_normalization=True + ) + mesh.write_data_to_vtk( + datasets={"flux": ref_data}, + filename=filename_without_norm, + volume_normalization=False + ) + + 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 + shape = (2, 2, 2) + np.testing.assert_allclose( + normalized_data, (ref_data.T.ravel() / cell_volume).reshape(shape)) + np.testing.assert_allclose( + unnormalized_data, ref_data.T.ravel().reshape(shape)) + + +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().reshape(2, 2, 2)) + + +@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"], +) +@pytest.mark.parametrize("volume_normalization", [False, True]) +def test_write_vtkhdf_rectilinear_multi_dimensional_data( + run_in_tmpdir, x_grid, y_grid, z_grid, data, expected, volume_normalization +): + """Test RectilinearMesh VTKHDF accepts 1D/2D/3D data with singleton padding. + + The grids have unit spacing, so every element volume is one and the + written values are the same with and without volume normalization. + """ + 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=volume_normalization, + ) + + 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] + + 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() + + shape = (2, 2, 2) + 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"] + + np.testing.assert_allclose( + root["CellData"]["flux"][()], data1.T.ravel().reshape(shape)) + np.testing.assert_allclose( + root["CellData"]["power"][()], data2.T.ravel().reshape(shape)) + np.testing.assert_allclose( + root["CellData"]["heating"][()], data3.T.ravel().reshape(shape)) + + +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"] + # a VTKHDF extent always has six entries, flat axes are padded to one + np.testing.assert_array_equal(root.attrs["WholeExtent"], + [0, 5, 0, 1, 0, 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).reshape(1, 1, 5)) + + grid = _read_vtkhdf(filename) + assert grid.GetNumberOfCells() == 5 + + +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"] + np.testing.assert_array_equal(root.attrs["WholeExtent"], + [0, 5, 0, 3, 0, 1]) + 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().reshape(1, 3, 5)) + + grid = _read_vtkhdf(filename) + assert grid.GetNumberOfCells() == 15 + + +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"""