Merge pull request #2464 from shimwell/moving_voxel_plot_code

moved voxel plot to api accessible function

docs/source/pythonapi/base.rst

@@ -192,6 +192,13 @@ Post-processing
    openmc.Track
    openmc.Tracks
 
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+   :template: myfunction.rst
+
+   openmc.voxel_to_vtk
+
 The following classes and functions are used for functional expansion reconstruction.
 
 .. autosummary::

openmc/plots.py

@@ -1,15 +1,18 @@
 from collections.abc import Iterable, Mapping
-from numbers import Real, Integral
+from numbers import Integral, Real
 from pathlib import Path
+from typing import Optional
 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 openmc.checkvalue import PathLike
 
+from ._xml import clean_indentation, get_elem_tuple, reorder_attributes
+from .mixin import IDManagerMixin
 
 _BASES = ['xy', 'xz', 'yz']
 
@@ -178,6 +181,78 @@ def _get_plot_image(plot, cwd):
     return Image(str(png_file))
 
 
+def voxel_to_vtk(voxel_file: PathLike, output: PathLike = 'plot.vti'):
+    """Converts a voxel HDF5 file to a VTK file
+
+    .. versionadded:: 0.13.4
+
+    Parameters
+    ----------
+    voxel_file : path-like
+        Path of the input h5 to convert
+    output : path-like
+        Path of the output vti file produced
+
+    Returns
+    -------
+    Path
+        Path of the .vti file produced
+    """
+
+    # imported vtk only if used as vtk is an option dependency
+    import vtk
+
+    _min_version = (2, 0)
+
+    # Read data from voxel file
+    with h5py.File(voxel_file, "r") as fh:
+        # 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
+        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(str(output))
+    writer.Write()
+
+    return output
+
+
 class PlotBase(IDManagerMixin):
     """
     Parameters
@@ -512,7 +587,7 @@ class Plot(PlotBase):
 
     @type.setter
     def type(self, plottype):
-        cv.check_value('plot type', plottype, ['slice', 'voxel', 'projection'])
+        cv.check_value('plot type', plottype, ['slice', 'voxel'])
         self._type = plottype
 
     @basis.setter
@@ -831,6 +906,45 @@ class Plot(PlotBase):
         # Return produced image
         return _get_plot_image(self, cwd)
 
+    def to_vtk(self, output: Optional[PathLike] = None,
+               openmc_exec: str = 'openmc', cwd: str = '.'):
+        """Render plot as an voxel image
+
+        This method runs OpenMC in plotting mode to produce a .vti file.
+
+        .. versionadded:: 0.13.4
+
+        Parameters
+        ----------
+        output : path-like
+            Path of the output .vti file produced
+        openmc_exec : str
+            Path to OpenMC executable
+        cwd : str, optional
+            Path to working directory to run in
+
+        Returns
+        -------
+        Path
+            Path of the .vti file produced
+
+        """
+        if self.type != 'voxel':
+            raise ValueError('Generating a VTK file only works for voxel plots')
+
+        # 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'
+        if output is None:
+            output = h5_voxel_file.with_suffix('.vti')
+
+        return voxel_to_vtk(h5_voxel_file, output)
+
 
 class ProjectionPlot(PlotBase):
     """Definition of a camera's view of OpenMC geometry
@@ -880,7 +994,7 @@ class ProjectionPlot(PlotBase):
     wireframe_domains : iterable of either Material or Cells
         If provided, the wireframe is only drawn around these.
         If color_by is by material, it must be a list of materials, else cells.
-    xs : dict 
+    xs : dict
         A mapping from cell/material IDs to floats. The floating point values
         are macroscopic cross sections influencing the volume rendering opacity
         of each geometric region. Zero corresponds to perfect transparency, and

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.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')
@@ -66,6 +69,35 @@ def myprojectionplot():
     return plot
 
 
+def test_voxel_plot(run_in_tmpdir):
+    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_vtk('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_vtk('another_test_voxel_plot.vti')
+
+    assert Path('h5_voxel_plot.h5').is_file()
+    assert Path('another_test_voxel_plot.vti').is_file()
+
+    slice_plot = openmc.Plot()
+    with pytest.raises(ValueError):
+        slice_plot.to_vtk('shimmy.vti')
+
+
 def test_attributes(myplot):
     assert myplot.name == 'myplot'