Auto origin and width for plots (#2492)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

openmc/bounding_box.py

@@ -27,6 +27,8 @@ class BoundingBox(tuple):
         The x, y, z coordinates of the upper right corner of the bounding box in [cm]
     volume : float
         The volume of the bounding box in [cm^3]
+    width : iterable of float
+        The width of the x, y and z axis in [cm]
     """
 
     def __new__(cls, lower_left: Iterable[float], upper_right: Iterable[float]):
@@ -55,3 +57,7 @@ class BoundingBox(tuple):
     @property
     def volume(self) -> float:
         return np.abs(np.prod(self[1] - self[0]))
+
+    @property
+    def width(self):
+        return self.upper_right - self.lower_left

openmc/universe.py

@@ -300,7 +300,7 @@ class Universe(UniverseBase):
     _default_legend_kwargs = {'bbox_to_anchor': (
         1.05, 1), 'loc': 2, 'borderaxespad': 0.0}
 
-    def plot(self, origin=(0., 0., 0.), width=(1., 1.), pixels=(200, 200),
+    def plot(self, origin=None, width=None, pixels=(200, 200),
              basis='xy', color_by='cell', colors=None, seed=None,
              openmc_exec='openmc', axes=None, legend=False,
              legend_kwargs=_default_legend_kwargs, outline=False,
@@ -309,10 +309,16 @@ class Universe(UniverseBase):
 
         Parameters
         ----------
-        origin : Iterable of float
-            Coordinates at the origin of the plot
-        width : Iterable of float
-            Width of the plot in each basis direction
+        origin : iterable of float
+            Coordinates at the origin of the plot, if left as None then the
+            universe.bounding_box.center will be used to attempt to
+            ascertain the origin. Defaults to (0, 0, 0) if the bounding_box
+            contains inf values
+        width : iterable of float
+            Width of the plot in each basis direction. If left as none then the
+            universe.bounding_box.width() will be used to attempt to
+            ascertain the plot width.  Defaults to (10, 10) if the bounding_box
+            contains inf values
         pixels : Iterable of int
             Number of pixels to use in each basis direction
         basis : {'xy', 'xz', 'yz'}
@@ -373,6 +379,22 @@ class Universe(UniverseBase):
         elif basis == 'xz':
             x, y = 0, 2
             xlabel, ylabel = 'x [cm]', 'z [cm]'
+
+        # checks to see if bounding box contains -inf or inf values
+        if True in np.isinf(self.bounding_box):
+            if origin is None:
+                origin = (0, 0, 0)
+            if width is None:
+                width = (10, 10)
+        else:
+            if origin is None:
+                origin = self.bounding_box.center
+            if width is None:
+                bb_width = self.bounding_box.width
+                x_width = bb_width['xyz'.index(basis[0])]
+                y_width = bb_width['xyz'.index(basis[1])]
+                width = (x_width, y_width)
+
         x_min = origin[x] - 0.5*width[0]
         x_max = origin[x] + 0.5*width[0]
         y_min = origin[y] - 0.5*width[1]

tests/unit_tests/test_universe.py

@@ -48,8 +48,9 @@ def test_bounding_box():
     assert_unbounded(u)
 
 
-def test_plot(run_in_tmpdir):
+def test_plot(run_in_tmpdir, sphere_model):
 
+    # model with -inf and inf in the bounding box
     pincell = openmc.examples.pwr_pin_cell()
     materials = pincell.materials
 
@@ -69,6 +70,22 @@ def test_plot(run_in_tmpdir):
             outline=True
         )
 
+    # model with no inf values in bounding box
+    m = sphere_model.materials[0]
+    univ = sphere_model.geometry.root_universe
+
+    colors = {m: 'limegreen'}
+
+    for basis in ('xy', 'yz', 'xz'):
+        univ.plot(
+            colors=colors,
+            color_by="cell",
+            legend=False,
+            pixels=(10, 10),
+            basis=basis,
+            outline=False
+        )
+
 
 def test_get_nuclides(uo2):
     c = openmc.Cell(fill=uo2)