Generate Region Plots Directly (openmc-dev#2895)

Co-authored-by: Paul Romano <[email protected]>

openmc/cell.py

@@ -594,7 +594,7 @@ def plot(self, *args, **kwargs):
 
                # Make water blue
                water = openmc.Cell(fill=h2o)
-               universe.plot(..., colors={water: (0., 0., 1.))
+               water.plot(colors={water: (0., 0., 1.)})
         seed : int
             Seed for the random number generator
         openmc_exec : str
@@ -619,8 +619,10 @@ def plot(self, *args, **kwargs):
 
         """
         # Create dummy universe but preserve used_ids
-        u = openmc.Universe(cells=[self], universe_id=openmc.Universe.next_id + 1)
+        next_id = openmc.Universe.next_id
+        u = openmc.Universe(cells=[self])
         openmc.Universe.used_ids.remove(u.id)
+        openmc.Universe.next_id = next_id
         return u.plot(*args, **kwargs)
 
     def create_xml_subelement(self, xml_element, memo=None):

openmc/region.py

@@ -1,9 +1,11 @@
 from abc import ABC, abstractmethod
 from collections.abc import MutableSequence
 from copy import deepcopy
+import warnings
 
 import numpy as np
 
+import openmc
 from .bounding_box import BoundingBox
 
 
@@ -334,6 +336,59 @@ def rotate(self, rotation, pivot=(0., 0., 0.), order='xyz', inplace=False,
         return type(self)(n.rotate(rotation, pivot=pivot, order=order,
                                    inplace=inplace, memo=memo) for n in self)
 
+    def plot(self, *args, **kwargs):
+        """Display a slice plot of the region.
+
+        .. versionadded:: 0.14.1
+
+        Parameters
+        ----------
+        origin : iterable of float
+            Coordinates at the origin of the plot. If left as None then the
+            bounding box center will be used to attempt to ascertain the origin.
+            Defaults to (0, 0, 0) if the bounding box is not finite
+        width : iterable of float
+            Width of the plot in each basis direction. If left as none then the
+            bounding box width will be used to attempt to ascertain the plot
+            width. Defaults to (10, 10) if the bounding box is not finite
+        pixels : Iterable of int or int
+            If iterable of ints provided, then this directly sets the number of
+            pixels to use in each basis direction. If int provided, then this
+            sets the total number of pixels in the plot and the number of pixels
+            in each basis direction is calculated from this total and the image
+            aspect ratio.
+        basis : {'xy', 'xz', 'yz'}
+            The basis directions for the plot
+        seed : int
+            Seed for the random number generator
+        openmc_exec : str
+            Path to OpenMC executable.
+        axes : matplotlib.Axes
+            Axes to draw to
+        outline : bool
+            Whether outlines between color boundaries should be drawn
+        axis_units : {'km', 'm', 'cm', 'mm'}
+            Units used on the plot axis
+        **kwargs
+            Keyword arguments passed to :func:`matplotlib.pyplot.imshow`
+
+        Returns
+        -------
+        matplotlib.axes.Axes
+            Axes containing resulting image
+
+        """
+        for key in ('color_by', 'colors', 'legend', 'legend_kwargs'):
+            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.plot(*args, **kwargs)
+
 
 class Intersection(Region, MutableSequence):
     r"""Intersection of two or more regions.

tests/unit_tests/test_region.py

@@ -241,3 +241,13 @@ def test_invalid_operands():
     with pytest.raises(ValueError, match='must be of type Region'):
         openmc.Complement(z)
 
+
+def test_plot():
+    # Create region and plot
+    region = -openmc.Sphere() & +openmc.XPlane()
+    c_before = openmc.Cell()
+    region.plot()
+
+    # Ensure that calling plot doesn't affect cell ID space
+    c_after = openmc.Cell()
+    assert c_after.id - 1 == c_before.id