added voxel plot

openmc/cell.py

@@ -605,6 +605,19 @@ def plot(self, *args, **kwargs):
         openmc.UniverseBase.next_id = next_id
         return u.plot(*args, **kwargs)
 
+    def voxel_plot(self, *args, **kwargs):
+        """Create a 3D voxel plot of the cell.
+
+        .. versionadded:: 0.15.2
+        """
+        # Create dummy universe but preserve used_ids
+        next_id = openmc.UniverseBase.next_id
+        u = openmc.Universe(cells=[self])
+        openmc.UniverseBase.used_ids.remove(u.id)
+        openmc.UniverseBase.next_id = next_id
+        return u.voxel_plot(*args, **kwargs)
+
+
     def create_xml_subelement(self, xml_element, memo=None):
         """Add the cell's xml representation to an incoming xml element
 

openmc/geometry.py

@@ -781,3 +781,37 @@ def plot(self, *args, **kwargs):
                 all_material_names.add(name)
 
         return model.plot(*args, **kwargs)
+
+    def voxel_plot(self, *args, **kwargs):
+        """Create a 3D voxel plot of the geometry.
+
+        .. versionadded:: 0.15.4
+        """
+        model = openmc.Model()
+        model.geometry = self
+        model.materials = self.get_all_materials().values()
+
+        # collect all the material names from the geometry
+        all_material_names = {m.name for m in model.materials if m.name is not None}
+
+        # makes a placeholder material for each material name if it isn't
+        # already present on the model. These materials are otherwise missing
+        # from the geometry and are needed for plotting.
+        for universe in model.geometry.get_all_universes().values():
+            if not isinstance(universe, openmc.DAGMCUniverse):
+                continue
+            for name in universe.material_names:
+                # if this name is already present in the model, skip it
+                # (this can happen if the same material name is used in multiple
+                # universes)
+                if name in all_material_names:
+                    continue
+                # if the material name is not present on the model,
+                # create a placeholder material with the same name
+                # and add it to the model
+                mat_dag = openmc.Material(name=name)
+                mat_dag.add_nuclide('H1', 1.0)
+                model.materials.append(mat_dag)
+                all_material_names.add(name)
+
+        return model.voxel_plot(*args, **kwargs)

openmc/model/model.py

@@ -1031,6 +1031,175 @@ def _set_plot_defaults(
 
         return origin, width, pixels
 
+    def _set_voxel_defaults(
+        self,
+        origin: Sequence[float] | None,
+        width: Sequence[float] | None,
+        pixels: int | Sequence[int],
+    ):
+        """Set default values for voxel plot parameters.
+
+        Parameters
+        ----------
+        origin : Sequence[float] or None
+            Origin (center) of the voxel plot
+        width : Sequence[float] or None
+            Width of the voxel plot in each dimension
+        pixels : int or Sequence[int]
+            Number of voxels in each direction
+
+        Returns
+        -------
+        tuple
+            (origin, width, pixels) with defaults applied
+        """
+        bb = self.bounding_box
+
+        # Check if bounding box contains inf values
+        if np.isinf(bb.extent['xyz']).any():
+            if origin is None:
+                origin = (0.0, 0.0, 0.0)
+            if width is None:
+                width = (10.0, 10.0, 10.0)
+        else:
+            if origin is None:
+                # Replace NaN values with 0.0
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore", RuntimeWarning)
+                    origin = np.nan_to_num(bb.center)
+            if width is None:
+                width = bb.width
+
+        # Convert single int pixels to 3D tuple with cubic voxels
+        if isinstance(pixels, int):
+            # Calculate cubic root to get voxels per side
+            voxels_per_side = (pixels ** (1/3))
+
+            # Distribute voxels proportionally to width in each dimension
+            # to maintain cubic voxels
+            pixels = tuple(
+                max(1, int(round(voxels_per_side * w / max(width))))
+                for w in width
+            )
+
+        return origin, width, pixels
+
+    def voxel_plot(
+        self,
+        origin: Sequence[float] | None = None,
+        width: Sequence[float] | None = None,
+        pixels: int | Sequence[int] = 64000,
+        color_by: str = 'cell',
+        colors: dict | None = None,
+        seed: int | None = None,
+        openmc_exec: PathLike = 'openmc',
+        output: PathLike = 'voxel_plot.vti',
+        show_overlaps: bool = False,
+        overlap_color: Sequence[int] | str | None = None,
+    ) -> Path:
+        """Create a 3D voxel plot of the model.
+
+        .. versionadded:: 0.15.4
+
+        Parameters
+        ----------
+        origin : Sequence[float], optional
+            Origin (center) of the plot (3 values). If unspecified, defaults to
+            the center of the bounding box if available, otherwise (0, 0, 0).
+        width : Sequence[float], optional
+            Width of the plot in each dimension (3 values). If unspecified,
+            defaults to the bounding box width if available, otherwise
+            (10, 10, 10).
+        pixels : int | Sequence[int], optional
+            If an iterable of ints is provided then this directly sets the
+            number of voxels in each direction (3 values). If a single int is
+            provided then this sets the total number of voxels in the plot and
+            the number of voxels in each direction is calculated to maintain
+            cubic voxels based on the width.
+        color_by : {'cell', 'material'}, optional
+            Indicate whether the plot should be colored by cell or by material.
+        colors : dict, optional
+            Dictionary indicating that certain cells/materials should be
+            displayed with a particular color.
+        seed : int, optional
+            Pseudorandom number seed for plot coloring
+        openmc_exec : PathLike, optional
+            Path to OpenMC executable. Defaults to 'openmc'.
+        output : PathLike, optional
+            Path to the output file. File extension determines format: '.h5'
+            for HDF5 voxel file, '.vti' for VTK image data file. Defaults to
+            'voxel_plot.vti'.
+        show_overlaps : bool, optional
+            Indicate whether or not overlapping regions are shown.
+        overlap_color : Sequence[int] or str, optional
+            Color to apply to overlapping regions.
+
+        Returns
+        -------
+        Path
+            Path to the generated voxel plot file (.h5 or .vti)
+
+        """
+        from openmc.plots import voxel_to_vtk
+        import shutil
+
+        # Set defaults using helper function
+        origin, width, pixels = self._set_voxel_defaults(origin, width, pixels)
+
+        # Validate output file extension
+        output_path = Path(output)
+        if output_path.suffix not in ('.h5', '.vti'):
+            raise ValueError(
+                f"Output file extension must be '.h5' or '.vti', "
+                f"got '{output_path.suffix}'"
+            )
+
+        with TemporaryDirectory() as tmpdir:
+            _plot_seed = self.settings.plot_seed
+            if seed is not None:
+                self.settings.plot_seed = seed
+
+            # Create voxel plot object
+            voxel = openmc.VoxelPlot()
+            voxel.origin = origin
+            voxel.width = width
+            voxel.pixels = pixels
+            voxel.color_by = color_by
+            voxel.show_overlaps = show_overlaps
+            if overlap_color is not None:
+                voxel.overlap_color = overlap_color
+            if colors is not None:
+                voxel.colors = colors
+
+            self.plots.append(voxel)
+
+            # Run OpenMC in geometry plotting mode
+            self.plot_geometry(False, cwd=tmpdir, openmc_exec=openmc_exec)
+
+            # Undo changes to model
+            self.plots.pop()
+            self.settings._plot_seed = _plot_seed
+
+            # Determine base filename
+            base_name = output_path.stem
+            h5_filename = f'{base_name}.h5'
+            h5_src = Path(tmpdir) / h5_filename
+
+            # Convert to VTK if .vti extension
+            if output_path.suffix == '.vti':
+                vti_filename = f'{base_name}.vti'
+
+                # Use voxel_to_vtk function from openmc.plots
+                vti_src = voxel_to_vtk(h5_src, Path(tmpdir) / vti_filename)
+
+                shutil.move(str(vti_src), output_path)
+
+                return output_path
+            else:  # .h5 extension
+                shutil.move(str(h5_src), output_path)
+
+                return output_path
+
     def id_map(
         self,
         origin: Sequence[float] | None = None,

openmc/region.py

@@ -361,6 +361,22 @@ def plot(self, *args, **kwargs):
         openmc.Cell.next_id = next_id
         return c.plot(*args, **kwargs)
 
+    def voxel_plot(self, *args, **kwargs):
+        """Create a 3D voxel plot of the region.
+
+        .. versionadded:: 0.15.2
+        """
+        for key in ('color_by', 'colors'):
+            if key in kwargs:
+                warnings.warn(f"The '{key}' argument is present but won't be applied in a region plot")
+
+        # Create cell while not perturbing use of autogenerated IDs
+        next_id = openmc.Cell.next_id
+        c = openmc.Cell(region=self)
+        openmc.Cell.used_ids.remove(c.id)
+        openmc.Cell.next_id = next_id
+        return c.voxel_plot(*args, **kwargs)
+
 
 class Intersection(Region, MutableSequence):
     r"""Intersection of two or more regions.

openmc/universe.py

@@ -339,6 +339,15 @@ def plot(self, *args, **kwargs):
         model.geometry = openmc.Geometry(self)
         return model.plot(*args, **kwargs)
 
+    def voxel_plot(self, *args, **kwargs):
+        """Create a 3D voxel plot of the universe.
+
+        .. versionadded:: 0.15.2
+        """
+        model = openmc.Model()
+        model.geometry = openmc.Geometry(self)
+        return model.voxel_plot(*args, **kwargs)
+
     def get_nuclides(self):
         """Returns all nuclides in the universe
 

tests/unit_tests/test_cell.py

@@ -384,3 +384,36 @@ def test_plot(run_in_tmpdir):
     # ensure that calling the plot method doesn't
     # affect the universe ID space
     assert u_before.id + 1 == u_after.id
+
+
+def test_voxel_plot(run_in_tmpdir):
+    """Test Cell.voxel_plot() method"""
+
+    # Create a material
+    mat = openmc.Material(material_id=1, name='test_mat')
+    mat.set_density('g/cm3', 1.0)
+    mat.add_element('H', 1.0)
+
+    # Create a simple cell with a sphere
+    sphere = openmc.Sphere(r=5.0, boundary_type='vacuum')
+    c = openmc.Cell(fill=mat, region=-sphere)
+
+    # Create a universe before the voxel_plot
+    u_before = openmc.Universe()
+
+    # Test VTI generation
+    vti_path = c.voxel_plot(pixels=1000, output='cell_test.vti')
+    assert vti_path.exists()
+    assert vti_path.suffix == '.vti'
+
+    # Test H5 generation
+    h5_path = c.voxel_plot(pixels=1000, output='cell_test.h5')
+    assert h5_path.exists()
+    assert h5_path.suffix == '.h5'
+
+    # Create a universe after the voxel_plot
+    u_after = openmc.Universe()
+
+    # Ensure that calling the voxel_plot method doesn't
+    # affect the universe ID space
+    assert u_before.id + 1 == u_after.id

tests/unit_tests/test_geometry.py

@@ -403,3 +403,45 @@ def test_redundant_surfaces():
     geom = openmc.Geometry([c3])
     redundant_surfs = geom.remove_redundant_surfaces()
     assert len(redundant_surfs) == 0
+
+
+def test_geometry_voxel_plot(run_in_tmpdir):
+    """Test Geometry.voxel_plot() method"""
+
+    # Create materials
+    mat1 = openmc.Material(material_id=1, name='material_1')
+    mat1.set_density('g/cm3', 1.0)
+    mat1.add_element('H', 1.0)
+
+    mat2 = openmc.Material(material_id=2, name='material_2')
+    mat2.set_density('g/cm3', 2.0)
+    mat2.add_element('Fe', 1.0)
+
+    # Create simple geometry with two spheres
+    inner_sphere = openmc.Sphere(r=5.0)
+    outer_sphere = openmc.Sphere(r=10.0, boundary_type='vacuum')
+
+    inner_cell = openmc.Cell(fill=mat1, region=-inner_sphere)
+    outer_cell = openmc.Cell(fill=mat2, region=+inner_sphere & -outer_sphere)
+
+    geometry = openmc.Geometry([inner_cell, outer_cell])
+
+    # Test 1: Generate VTI file
+    vti_path = geometry.voxel_plot(pixels=1000, output='geom_test.vti')
+    assert vti_path.exists()
+    assert vti_path.suffix == '.vti'
+
+    # Test 2: Generate H5 file
+    h5_path = geometry.voxel_plot(pixels=1000, output='geom_test.h5')
+    assert h5_path.exists()
+    assert h5_path.suffix == '.h5'
+
+    # Test 3: Test with explicit pixel dimensions
+    explicit_path = geometry.voxel_plot(pixels=(10, 10, 10), output='geom_explicit.vti')
+    assert explicit_path.exists()
+
+    # Test 4: Test default output
+    default_path = geometry.voxel_plot(pixels=1000)
+    assert default_path.exists()
+    assert default_path.name == 'voxel_plot.vti'
+