diff --git a/openmc/mesh.py b/openmc/mesh.py index 3c3c0a1ac5d..3d8da0a39ec 100644 --- a/openmc/mesh.py +++ b/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,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 @@ -728,20 +742,26 @@ def write_data_to_vtk(self, # maintain a list of the datasets as added to the VTK arrays to # ensure they persist in memory until the file is written datasets_out = [] + # Regular/Rectilinear meshes store data in C (ijk) order which + # matches VTK's expected ordering directly — no transpose needed. + # Curvilinear meshes (Cylindrical, Spherical) require a transpose + # to convert from C ordering to the Fortran (kji) ordering that VTK + # expects. + # TODO: update to "C" ordering throughout + needs_transpose = not isinstance( + self, (RegularMesh, RectilinearMesh)) for label, dataset in datasets.items(): dataset = self._reshape_vtk_dataset(dataset) self._check_vtk_dataset(label, dataset) - # If the array data is 3D, assume is in C ordering and transpose - # before flattening to match the ordering expected by the VTK - # array based on the way mesh indices are ordered in the Python - # API - # TODO: update to "C" ordering throughout if dataset.ndim == 3: - dataset = dataset.T.ravel() + dataset = dataset.T.ravel() if needs_transpose else dataset.ravel() datasets_out.append(dataset) if volume_normalization: - dataset /= self.volumes.T.ravel() + if needs_transpose: + dataset /= self.volumes.T.ravel() + else: + dataset /= self.volumes.ravel() dataset_array = vtk.vtkDoubleArray() dataset_array.SetName(label) @@ -929,12 +949,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 +1572,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 +1880,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 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) + _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 +2402,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 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) + _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 +2819,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 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) + _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 +3319,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] diff --git a/tests/unit_tests/test_mesh.py b/tests/unit_tests/test_mesh.py index 0b28bdfbe44..8d9568218b0 100644 --- a/tests/unit_tests/test_mesh.py +++ b/tests/unit_tests/test_mesh.py @@ -995,6 +995,415 @@ 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] + + # Create data with known values + ref_data = np.ones(mesh.dimension) * 100.0 + 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.ravel() / cell_volume + np.testing.assert_allclose(normalized_data, expected_normalized) + np.testing.assert_allclose(unnormalized_data, ref_data.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] + + rng = np.random.default_rng(42) + ref_data = rng.random(mesh.dimension) + filename = "test_1d_mesh.vtkhdf" + mesh.write_data_to_vtk(datasets={"data": ref_data}, filename=filename, + volume_normalization=False) + + assert Path(filename).exists() + + with h5py.File(filename, "r") as f: + root = f["VTKHDF"] + # Verify 1D mesh was written + dims = root["Dimensions"][()] + assert len(dims) >= 1 + assert "CellData" in root + + +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] + + rng = np.random.default_rng(42) + ref_data = rng.random(mesh.dimension) + filename = "test_2d_mesh.vtkhdf" + mesh.write_data_to_vtk(datasets={"data": ref_data}, filename=filename, + volume_normalization=False) + + assert Path(filename).exists() + + with h5py.File(filename, "r") as f: + root = f["VTKHDF"] + # Verify 2D mesh was written + dims = root["Dimensions"][()] + assert len(dims) == 2 + assert "CellData" in root + + +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.ravel()) def test_rectilinear_mesh_get_indices_at_coords(): """Test get_indices_at_coords method for RectilinearMesh"""