diff --git a/openmc/mesh.py b/openmc/mesh.py index 7f7b078795..2a53c8a91e 100644 --- a/openmc/mesh.py +++ b/openmc/mesh.py @@ -193,6 +193,77 @@ class StructuredMesh(MeshBase): s1 = (slice(1, None),)*ndim + (slice(None),) return (vertices[s0] + vertices[s1]) / 2 + def write_data_to_vtk(self, points, filename, datasets, volume_normalization=True): + """Creates a VTK object of the mesh + + Parameters + ---------- + points : list or np.array + List of (X,Y,Z) tuples. + filename : str + Name of the VTK file to write. + datasets : dict + Dictionary whose keys are the data labels + and values are the data sets. + volume_normalization : bool, optional + Whether or not to normalize the data by + the volume of the mesh elements. + + Raises + ------ + RuntimeError + When the size of a dataset doesn't match the number of cells + + Returns + ------- + vtk.vtkStructuredGrid + the VTK object + """ + + import vtk + from vtk.util import numpy_support as nps + + # check that the data sets are appropriately sized + errmsg = "The size of the dataset {} should be equal to the number of cells" + for label, dataset in datasets.items(): + if isinstance(dataset, np.ndarray): + if not dataset.size == self.dimension[0] * self.dimension[1]* self.dimension[2]: + raise RuntimeError(errmsg.format(label)) + else: + if len(dataset) == self.dimension[0] * self.dimension[1]* self.dimension[2]: + raise RuntimeError(errmsg.format(label)) + cv.check_type('label', label, str) + + vtk_grid = vtk.vtkStructuredGrid() + + vtk_grid.SetDimensions(*self.dimension) + + vtkPts = vtk.vtkPoints() + vtkPts.SetData(nps.numpy_to_vtk(points, deep=True)) + vtk_grid.SetPoints(vtkPts) + + # create VTK arrays for each of + # the data sets + for label, dataset in datasets.items(): + dataset = np.asarray(dataset).flatten() + + if volume_normalization: + dataset /= self.volumes.flatten() + + dataset_array = vtk.vtkDoubleArray() + dataset_array.SetName(label) + dataset_array.SetArray(nps.numpy_to_vtk(dataset), + dataset.size, + True) + vtk_grid.GetCellData().AddArray(dataset_array) + + # write the .vtk file + writer = vtk.vtkStructuredGridWriter() + writer.SetFileName(str(filename)) + writer.SetInputData(vtk_grid) + writer.Write() + + return vtk_grid class RegularMesh(StructuredMesh): """A regular Cartesian mesh in one, two, or three dimensions @@ -621,7 +692,42 @@ class RegularMesh(StructuredMesh): root_cell.fill = lattice return root_cell, cells + + def write_data_to_vtk(self, filename, datasets, volume_normalization=True): + """Creates a VTK object of the mesh + Parameters + ---------- + filename : str or pathlib.Path + Name of the VTK file to write. + datasets : dict + Dictionary whose keys are the data labels + and values are the data sets. + volume_normalization : bool, optional + Whether or not to normalize the data by + the volume of the mesh elements. + Defaults to True. + + Returns + ------- + vtk.vtkStructuredGrid + the VTK object + """ + + ll, ur = self.lower_left, self.upper_right + x_vals = np.linspace(ll[0], ur[0], num=self.dimension[0] + 1) + y_vals = np.linspace(ll[1], ur[1], num=self.dimension[1] + 1) + z_vals = np.linspace(ll[2], ur[2], num=self.dimension[2] + 1) + + # create points + pts_cartesian = np.array([[x, y, z] for z in z_vals for y in y_vals for x in x_vals]) + + return super().write_data_to_vtk( + points=pts_cartesian, + filename=filename, + datasets=datasets, + volume_normalization=volume_normalization + ) def Mesh(*args, **kwargs): warnings.warn("Mesh has been renamed RegularMesh. Future versions of " @@ -821,6 +927,36 @@ class RectilinearMesh(StructuredMesh): return element + def write_data_to_vtk(self, filename, datasets, volume_normalization=True): + """Creates a VTK object of the mesh + + Parameters + ---------- + filename : str or pathlib.Path + Name of the VTK file to write. + datasets : dict + Dictionary whose keys are the data labels + and values are the data sets. + volume_normalization : bool, optional + Whether or not to normalize the data by + the volume of the mesh elements. + Defaults to True. + + Returns + ------- + vtk.vtkStructuredGrid + the VTK object + """ + # create points + pts_cartesian = np.array([[x, y, z] for z in self.z_grid for y in self.y_grid for x in self.x_grid]) + + return super().write_data_to_vtk( + points=pts_cartesian, + filename=filename, + datasets=datasets, + volume_normalization=volume_normalization + ) + class CylindricalMesh(StructuredMesh): """A 3D cylindrical mesh @@ -1012,6 +1148,46 @@ class CylindricalMesh(StructuredMesh): return np.multiply.outer(np.outer(V_r, V_p), V_z) + def write_data_to_vtk(self, filename, datasets, volume_normalization=True): + """Creates a VTK object of the mesh + + Parameters + ---------- + filename : str or pathlib.Path + Name of the VTK file to write. + datasets : dict + Dictionary whose keys are the data labels + and values are the data sets. + volume_normalization : bool, optional + Whether or not to normalize the data by + the volume of the mesh elements. + Defaults to True. + + Returns + ------- + vtk.vtkStructuredGrid + the VTK object + """ + # create points + pts_cylindrical = np.array( + [ + [r, phi, z] + for z in self.z_grid + for phi in self.phi_grid + for r in self.r_grid + ] + ) + pts_cartesian = np.copy(pts_cylindrical) + r, phi = pts_cylindrical[:, 0], pts_cylindrical[:, 1] + pts_cartesian[:, 0] = r * np.cos(phi) + pts_cartesian[:, 1] = r * np.sin(phi) + + return super().write_data_to_vtk( + points=pts_cartesian, + filename=filename, + datasets=datasets, + volume_normalization=volume_normalization + ) class SphericalMesh(StructuredMesh): """A 3D spherical mesh @@ -1204,6 +1380,49 @@ class SphericalMesh(StructuredMesh): return np.multiply.outer(np.outer(V_r, V_t), V_p) + def write_data_to_vtk(self, filename, datasets, volume_normalization=True): + """Creates a VTK object of the mesh + + Parameters + ---------- + filename : str or pathlib.Path + Name of the VTK file to write. + datasets : dict + Dictionary whose keys are the data labels + and values are the data sets. + volume_normalization : bool, optional + Whether or not to normalize the data by + the volume of the mesh elements. + Defaults to True. + + Returns + ------- + vtk.vtkStructuredGrid + the VTK object + """ + + # create points + pts_spherical = np.array( + [ + [r, theta, phi] + for phi in self.phi_grid + for theta in self.theta_grid + for r in self.r_grid + ] + ) + pts_cartesian = np.copy(pts_spherical) + r, theta, phi = pts_spherical[:, 0], pts_spherical[:, 1], pts_spherical[:, 2] + pts_cartesian[:, 0] = r * np.sin(phi) * np.cos(theta) + pts_cartesian[:, 1] = r * np.sin(phi) * np.sin(theta) + pts_cartesian[:, 2] = r * np.cos(phi) + + return super().write_data_to_vtk( + points=pts_cartesian, + filename=filename, + datasets=datasets, + volume_normalization=volume_normalization + ) + class UnstructuredMesh(MeshBase): """A 3D unstructured mesh @@ -1373,7 +1592,7 @@ class UnstructuredMesh(MeshBase): Parameters ---------- - filename : str + filename : str or pathlib.Path Name of the VTK file to write. datasets : dict Dictionary whose keys are the data labels @@ -1381,6 +1600,11 @@ class UnstructuredMesh(MeshBase): volume_normalization : bool Whether or not to normalize the data by the volume of the mesh elements + + Raises + ------ + RuntimeError + when the size of a dataset doesn't match the number of cells """ import vtk @@ -1397,11 +1621,14 @@ class UnstructuredMesh(MeshBase): " mesh information from a statepoint file.") # check that the data sets are appropriately sized + errmsg = "The size of the dataset {} should be equal to the number of cells" for label, dataset in datasets.items(): if isinstance(dataset, np.ndarray): - assert dataset.size == self.n_elements + if not dataset.size == self.dimension[0] * self.dimension[1]* self.dimension[2]: + raise RuntimeError(errmsg.format(label)) else: - assert len(dataset) == self.n_elements + if len(dataset) == self.dimension[0] * self.dimension[1]* self.dimension[2]: + raise RuntimeError(errmsg.format(label)) cv.check_type('label', label, str) # create data arrays for the cells/points diff --git a/tests/unit_tests/test_mesh_to_vtk.py b/tests/unit_tests/test_mesh_to_vtk.py new file mode 100644 index 0000000000..9247540a09 --- /dev/null +++ b/tests/unit_tests/test_mesh_to_vtk.py @@ -0,0 +1,76 @@ +import numpy as np +from pathlib import Path +import pytest + +vtk = pytest.importorskip("vtk") +from vtk.util import numpy_support as nps + +import openmc + + +regular_mesh = openmc.RegularMesh() +regular_mesh.lower_left = (0, 0, 0) +regular_mesh.upper_right = (1, 1, 1) +regular_mesh.dimension = [30, 20, 10] + +rectilinear_mesh = openmc.RectilinearMesh() +rectilinear_mesh.x_grid = np.linspace(1, 2, num=30) +rectilinear_mesh.y_grid = np.linspace(1, 2, num=30) +rectilinear_mesh.z_grid = np.linspace(1, 2, num=30) + +cylinder_mesh = openmc.CylindricalMesh() +cylinder_mesh.r_grid = np.linspace(1, 2, num=30) +cylinder_mesh.phi_grid = np.linspace(0, np.pi, num=50) +cylinder_mesh.z_grid = np.linspace(0, 1, num=30) + +spherical_mesh = openmc.SphericalMesh() +spherical_mesh.r_grid = np.linspace(1, 2, num=30) +spherical_mesh.phi_grid = np.linspace(0, np.pi, num=50) +spherical_mesh.theta_grid = np.linspace(0, np.pi / 2, num=30) + +@pytest.mark.parametrize("mesh", [cylinder_mesh, regular_mesh, rectilinear_mesh, spherical_mesh]) +def test_write_data_to_vtk(mesh, tmpdir): + # BUILD + filename = Path(tmpdir) / "out.vtk" + + data = np.random.random(mesh.dimension[0]*mesh.dimension[1]*mesh.dimension[2]) + + # RUN + mesh.write_data_to_vtk(filename=filename, datasets={"label1": data, "label2": data}) + + # TEST + assert filename.is_file() + + # read file + reader = vtk.vtkStructuredGridReader() + reader.SetFileName(filename) + reader.Update() + + # check name of datasets + vtk_grid = reader.GetOutput() + array1 = vtk_grid.GetCellData().GetArray(0) + array2 = vtk_grid.GetCellData().GetArray(1) + + assert array1.GetName() == "label1" + assert array2.GetName() == "label2" + + # check size of datasets + assert nps.vtk_to_numpy(array1).size == data.size + assert nps.vtk_to_numpy(array2).size == data.size + +@pytest.mark.parametrize("mesh", [cylinder_mesh, regular_mesh, rectilinear_mesh, spherical_mesh]) +def test_write_data_to_vtk_size_mismatch(mesh): + """Checks that an error is raised when the size of the dataset + doesn't match the mesh number of cells + + Parameters + ---------- + mesh : openmc.StructuredMesh + The mesh to test + """ + right_size = mesh.dimension[0]*mesh.dimension[1]*mesh.dimension[2] + data = np.random.random(right_size + 1) + + expected_error_msg = "The size of the dataset label should be equal to the number of cells" + with pytest.raises(RuntimeError, match=expected_error_msg): + mesh.write_data_to_vtk(filename="out.vtk", datasets={"label": data}) \ No newline at end of file diff --git a/tests/unit_tests/test_tallies.py b/tests/unit_tests/test_tallies.py index aeeb0612d5..bfff741a97 100644 --- a/tests/unit_tests/test_tallies.py +++ b/tests/unit_tests/test_tallies.py @@ -1,4 +1,5 @@ import numpy as np + import openmc