Merge pull request #2464 from shimwell/moving_voxel_plot_code

moved voxel plot to api accessible function

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