moved voxel plot to api accessible function

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

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}...")

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