diff --git a/openmc/plots.py b/openmc/plots.py index c3cd88da18..573bc80eb6 100644 --- a/openmc/plots.py +++ b/openmc/plots.py @@ -1,15 +1,16 @@ from collections.abc import Iterable, Mapping -from numbers import Real, Integral +from numbers import Integral, Real from pathlib import Path from xml.etree import ElementTree as ET +import h5py import numpy as np import openmc import openmc.checkvalue as cv -from ._xml import clean_indentation, reorder_attributes, get_elem_tuple -from .mixin import IDManagerMixin +from ._xml import clean_indentation, get_elem_tuple, reorder_attributes +from .mixin import IDManagerMixin _BASES = ['xy', 'xz', 'yz'] @@ -178,6 +179,78 @@ def _get_plot_image(plot, cwd): return Image(str(png_file)) +def _voxel_to_vtk(voxel_file: str, output: str = 'plot.vti'): + """Converts a voxel HDF5 file to a VTK file') + + Parameters + ---------- + voxel_file : str + filename of the input h5 to convert + output : str + filename of the output vti file produced + + Returns + ------- + str + Filename of the vti file produced + """ + + # imported vkt only if used as vtk is an option dependency + import vtk + + _min_version = (2, 0) + + # Read data from voxel file + fh = h5py.File(voxel_file, "r") + + # check version + version = tuple(fh.attrs["version"]) + if version < _min_version: + old_version = ".".join(map(str, version)) + min_version = ".".join(map(str, _min_version)) + err_msg = ( + f"This voxel file's version is {old_version}. This function only " + f" supports voxel files with version {min_version} or higher. " + "Please generate a new voxel file using a newer version of OpenMC." + ) + raise ValueError(err_msg) + + dimension = fh.attrs["num_voxels"] + width = fh.attrs["voxel_width"] + lower_left = fh.attrs["lower_left"] + + nx, ny, nz = dimension + + grid = vtk.vtkImageData() + grid.SetDimensions(nx + 1, ny + 1, nz + 1) + grid.SetOrigin(*lower_left) + grid.SetSpacing(*width) + + # transpose data from OpenMC ordering (zyx) to VTK ordering (xyz) + # and flatten to 1-D array + print("Reading and translating data...") + h5data = fh["data"][...] + + data = vtk.vtkIntArray() + data.SetName("id") + # set the array using the h5data array + data.SetArray(h5data, h5data.size, True) + # add data to image grid + grid.GetCellData().AddArray(data) + + writer = vtk.vtkXMLImageDataWriter() + if vtk.vtkVersion.GetVTKMajorVersion() > 5: + writer.SetInputData(grid) + else: + writer.SetInput(grid) + if not output.endswith(".vti"): + output += ".vti" + writer.SetFileName(output) + writer.Write() + + return output + + class PlotBase(IDManagerMixin): """ Parameters @@ -831,6 +904,37 @@ class Plot(PlotBase): # Return produced image return _get_plot_image(self, cwd) + def to_voxel_file(self, output='plot.vti', openmc_exec='openmc', cwd='.'): + """Render plot as an voxel image + + This method runs OpenMC in plotting mode to produce a .vti file. + + Parameters + ---------- + output : str + filename of the output vti file produced + openmc_exec : str + Path to OpenMC executable + cwd : str, optional + Path to working directory to run in + + Returns + ------- + str + Filename of the vti file produced + + """ + # Create plots.xml + Plots([self]).export_to_xml(cwd) + + # Run OpenMC in geometry plotting mode and produces a h5 file + openmc.plot_geometry(False, openmc_exec, cwd) + + stem = self.filename if self.filename is not None else f'plot_{self.id}' + h5_voxel_file = Path(cwd) / f'{stem}.h5' + + return _voxel_to_vtk(h5_voxel_file, output) + class ProjectionPlot(PlotBase): """Definition of a camera's view of OpenMC geometry diff --git a/scripts/openmc-voxel-to-vtk b/scripts/openmc-voxel-to-vtk index 01b4808d2b..16789cb04e 100755 --- a/scripts/openmc-voxel-to-vtk +++ b/scripts/openmc-voxel-to-vtk @@ -2,63 +2,7 @@ from argparse import ArgumentParser -import h5py -import vtk - -_min_version = (2, 0) - - -def voxel_to_vtk(voxel_file: str, output: str): - - # Read data from voxel file - fh = h5py.File(voxel_file, "r") - - # check version - version = tuple(fh.attrs["version"]) - if version < _min_version: - old_version = ".".join(map(str, version)) - min_version = ".".join(map(str, _min_version)) - err_msg = ( - "This voxel file's version is {}. This script " - "only supports voxel files with version {} or " - "higher. Please generate a new voxel file using " - "a newer version of OpenMC.".format(old_version, min_version) - ) - raise ValueError(err_msg) - - dimension = fh.attrs["num_voxels"] - width = fh.attrs["voxel_width"] - lower_left = fh.attrs["lower_left"] - - nx, ny, nz = dimension - - grid = vtk.vtkImageData() - grid.SetDimensions(nx + 1, ny + 1, nz + 1) - grid.SetOrigin(*lower_left) - grid.SetSpacing(*width) - - # transpose data from OpenMC ordering (zyx) to VTK ordering (xyz) - # and flatten to 1-D array - print("Reading and translating data...") - h5data = fh["data"][...] - - data = vtk.vtkIntArray() - data.SetName("id") - # set the array using the h5data array - data.SetArray(h5data, h5data.size, True) - # add data to image grid - grid.GetCellData().AddArray(data) - - writer = vtk.vtkXMLImageDataWriter() - if vtk.vtkVersion.GetVTKMajorVersion() > 5: - writer.SetInputData(grid) - else: - writer.SetInput(grid) - if not output.endswith(".vti"): - output += ".vti" - writer.SetFileName(output) - print("Writing VTK file {}...".format(output)) - writer.Write() +import openmc if __name__ == "__main__": @@ -73,4 +17,6 @@ if __name__ == "__main__": help="Path to output VTK file.", ) args = parser.parse_args() - voxel_to_vtk(args.voxel_file, args.output) + print("Reading and translating data...") + openmc.plots._voxel_to_vtk(args.voxel_file, args.output) + print(f"Written VTK file {args.output}...") diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py index 183341c26e..b556be361e 100644 --- a/tests/unit_tests/test_plots.py +++ b/tests/unit_tests/test_plots.py @@ -1,3 +1,5 @@ +from pathlib import Path + import openmc import openmc.examples import pytest @@ -37,6 +39,7 @@ def myplot(): } return plot + @pytest.fixture(scope='module') def myprojectionplot(): plot = openmc.ProjectionPlot(name='myprojectionplot') @@ -59,26 +62,56 @@ def myprojectionplot(): plot.overlap_color = (255, 211, 0) plot.overlap_color = 'yellow' - + plot.wireframe_thickness = 2 plot.level = 1 return plot + +def test_voxel_plot(): + surf1 = openmc.Sphere(r=500, boundary_type='vacuum') + cell1 = openmc.Cell(region=-surf1) + geometry = openmc.Geometry([cell1]) + geometry.export_to_xml() + materials = openmc.Materials() + materials.export_to_xml() + vox_plot = openmc.Plot() + vox_plot.type = 'voxel' + vox_plot.id = 12 + vox_plot.width = (1500., 1500., 1500.) + vox_plot.pixels = (200, 200, 200) + vox_plot.color_by = 'cell' + vox_plot.to_voxel_file(output='test_voxel_plot.vti') + + assert Path('plot_12.h5').is_file() + assert Path('test_voxel_plot.vti').is_file() + + vox_plot.filename = 'h5_voxel_plot' + vox_plot.to_voxel_file(output='another_test_voxel_plot.vti') + + assert Path('h5_voxel_plot.h5').is_file() + assert Path('another_test_voxel_plot.vti').is_file() + + def test_attributes(myplot): assert myplot.name == 'myplot' + def test_attributes_proj(myprojectionplot): assert myprojectionplot.name == 'myprojectionplot' + def test_repr(myplot): r = repr(myplot) assert isinstance(r, str) + def test_repr_proj(myprojectionplot): r = repr(myprojectionplot) assert isinstance(r, str) + def test_from_geometry(): width = 25. s = openmc.Sphere(r=width/2, boundary_type='vacuum') @@ -154,7 +187,7 @@ def test_plots(run_in_tmpdir): plots.export_to_xml() # from_xml - new_plots = openmc.Plots.from_xml() + openmc.Plots.from_xml() assert len(plots) assert plots[0].origin == p1.origin assert plots[0].colors == p1.colors