Fixing crash when calling Geometry.plot when DAGMCUniverse in geometry (#3455)

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

openmc/dagmc.py

@@ -571,16 +571,7 @@ class DAGMCUniverse(openmc.UniverseBase):
     def plot(self, *args, **kwargs):
         """Display a slice plot of the DAGMCUniverse.
         """
-        model = openmc.Model()
-        model.geometry = openmc.Geometry(self)
-
-        for mat_name in self.material_names:
-            material = openmc.Material(name=mat_name)
-            # Placeholder nuclide to ensure material is not empty
-            material.add_nuclide('H1', 1.0)
-            model.materials.append(material)
-
-        return model.plot(*args, **kwargs)
+        return openmc.Geometry(self).plot(*args, **kwargs)
 
 
 class DAGMCCell(openmc.Cell):

openmc/geometry.py

@@ -754,4 +754,17 @@ class Geometry:
 
         .. versionadded:: 0.14.0
         """
-        return self.root_universe.plot(*args, **kwargs)
+        model = openmc.Model()
+        model.geometry = self
+        model.materials = self.get_all_materials().values()
+
+        # Add placeholder materials for DAGMCUniverses
+        universes = self.get_all_universes()
+        for universe in universes.values():
+            if isinstance(universe, openmc.DAGMCUniverse):
+                for name in universe.material_names:
+                    mat_dag = openmc.Material(name=name)
+                    mat_dag.add_nuclide('H1', 1.0)
+                    model.materials.append(mat_dag)
+
+        return model.plot(*args, **kwargs)

openmc/model/model.py

@@ -986,7 +986,16 @@ class Model:
         y_min = (origin[y] - 0.5*width[1]) * axis_scaling_factor[axis_units]
         y_max = (origin[y] + 0.5*width[1]) * axis_scaling_factor[axis_units]
 
+        # Determine whether any materials contains macroscopic data and if so,
+        # set energy mode accordingly
+        _energy_mode = self.settings._energy_mode
+        for mat in self.geometry.get_all_materials().values():
+            if mat._macroscopic is not None:
+                self.settings.energy_mode = 'multi-group'
+                break
+
         with TemporaryDirectory() as tmpdir:
+            _plot_seed = self.settings.plot_seed
             if seed is not None:
                 self.settings.plot_seed = seed
 
@@ -1007,6 +1016,11 @@ class Model:
             # 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
+            self.settings._energy_mode = _energy_mode
+
             # Read image from file
             img_path = Path(tmpdir) / f'plot_{plot.id}.png'
             if not img_path.is_file():

openmc/universe.py

@@ -337,14 +337,6 @@ class UniverseBase(ABC, IDManagerMixin):
         """
         model = openmc.Model()
         model.geometry = openmc.Geometry(self)
-
-        # Determine whether any materials contains macroscopic data and if
-        # so, set energy mode accordingly
-        for mat in self.get_all_materials().values():
-            if mat._macroscopic is not None:
-                model.settings.energy_mode = 'multi-group'
-                break
-
         return model.plot(*args, **kwargs)
 
     def get_nuclides(self):

tests/unit_tests/dagmc/test_plot.py

@@ -38,3 +38,24 @@ def test_plotting_dagmc_universe(request):
 
     dag_universe = openmc.DAGMCUniverse(request.path.parent / 'dagmc.h5m')
     dag_universe.plot()
+
+
+def test_plotting_geometry_filled_with_dagmc_universe(request):
+    """Test plotting a geometry with OpenMC. This is an edge case when plotting
+    geometry as often geometry objects don't include a DAGMCUniverse. The
+    inclusion of a DAGMCUniverse requires special handling for the materials."""
+
+    dag_universe = openmc.DAGMCUniverse(request.path.parent / 'dagmc.h5m', auto_geom_ids=True)
+
+    sphere1 = openmc.Sphere(r=50.0)
+    sphere2 = openmc.Sphere(r=60.0, boundary_type='vacuum')
+
+    # Adding a material to the CSG Universe to check all materials are accounted for
+    csg_material = openmc.Material(name='csg_material')
+    csg_material.add_nuclide("H1", 1.0)
+
+    cell1 = openmc.Cell(fill=dag_universe, region=-sphere1)
+    cell2 = openmc.Cell(fill=csg_material, region=+sphere1 & -sphere2)
+
+    geometry = openmc.Geometry([cell1, cell2])
+    geometry.plot()