Adding source option to plot (#2863)

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

openmc/geometry.py

@@ -7,8 +7,6 @@ from pathlib import Path
 import warnings
 import lxml.etree as ET
 
-import numpy as np
-
 import openmc
 import openmc._xml as xml
 from .checkvalue import check_type, check_less_than, check_greater_than, PathLike
@@ -67,7 +65,7 @@ class Geometry:
         self._root_universe = root_universe
 
     @property
-    def bounding_box(self) -> np.ndarray:
+    def bounding_box(self) -> openmc.BoundingBox:
         return self.root_universe.bounding_box
 
     @property
@@ -800,6 +798,7 @@ class Geometry:
             Units used on the plot axis
         **kwargs
             Keyword arguments passed to :func:`matplotlib.pyplot.imshow`
+
         Returns
         -------
         matplotlib.axes.Axes

openmc/lib/core.py

@@ -474,8 +474,11 @@ def run(output=True):
         _dll.openmc_run()
 
 
-def sample_external_source(n_samples=1, prn_seed=None):
-    """Sample external source
+def sample_external_source(
+        n_samples: int = 1000,
+        prn_seed: int | None = None
+) -> list[openmc.SourceParticle]:
+    """Sample external source and return source particles.
 
     .. versionadded:: 0.13.1
 
@@ -490,7 +493,7 @@ def sample_external_source(n_samples=1, prn_seed=None):
     Returns
     -------
     list of openmc.SourceParticle
-        List of samples source particles
+        List of sampled source particles
 
     """
     if n_samples <= 0:

openmc/model/model.py

@@ -9,6 +9,7 @@ import warnings
 
 import h5py
 import lxml.etree as ET
+import numpy as np
 
 import openmc
 import openmc._xml as xml
@@ -793,6 +794,121 @@ class Model:
                             openmc.lib.materials[domain_id].volume = \
                                 vol_calc.volumes[domain_id].n
 
+    def plot(
+        self,
+        n_samples: int | None = None,
+        plane_tolerance: float = 1.,
+        source_kwargs: dict | None = None,
+        **kwargs,
+    ):
+        """Display a slice plot of the geometry.
+
+        .. versionadded:: 0.15.1
+
+        Parameters
+        ----------
+        n_samples : dict
+            The number of source particles to sample and add to plot. Defaults
+            to None which doesn't plot any particles on the plot.
+        plane_tolerance: float
+            When plotting a plane the source locations within the plane +/-
+            the plane_tolerance will be included and those outside of the
+            plane_tolerance will not be shown
+        source_kwargs : dict
+            Keyword arguments passed to :func:`matplotlib.pyplot.scatter`.
+        **kwargs
+            Keyword arguments passed to :func:`openmc.Universe.plot`
+
+        Returns
+        -------
+        matplotlib.axes.Axes
+            Axes containing resulting image
+        """
+
+        check_type('n_samples', n_samples, int | None)
+        check_type('plane_tolerance', plane_tolerance, float)
+        if source_kwargs is None:
+            source_kwargs = {}
+        source_kwargs.setdefault('marker', 'x')
+
+        ax = self.geometry.plot(**kwargs)
+        if n_samples:
+            # Sample external source particles
+            particles = self.sample_external_source(n_samples)
+
+            # Determine plotting parameters and bounding box of geometry
+            bbox = self.geometry.bounding_box
+            origin = kwargs.get('origin', None)
+            basis = kwargs.get('basis', 'xy')
+            indices = {'xy': (0, 1, 2), 'xz': (0, 2, 1), 'yz': (1, 2, 0)}[basis]
+
+            # Infer origin if not provided
+            if np.isinf(bbox.extent[basis]).any():
+                if origin is None:
+                    origin = (0, 0, 0)
+            else:
+                if origin is None:
+                    # if nan values in the bbox.center they get replaced with 0.0
+                    # this happens when the bounding_box contains inf values
+                    with warnings.catch_warnings():
+                        warnings.simplefilter("ignore", RuntimeWarning)
+                        origin = np.nan_to_num(bbox.center)
+
+            slice_index = indices[2]
+            slice_value = origin[slice_index]
+
+            xs = []
+            ys = []
+            tol = plane_tolerance
+            for particle in particles:
+                if (slice_value - tol < particle.r[slice_index] < slice_value + tol):
+                    xs.append(particle.r[indices[0]])
+                    ys.append(particle.r[indices[1]])
+            ax.scatter(xs, ys, **source_kwargs)
+        return ax
+
+    def sample_external_source(
+            self,
+            n_samples: int = 1000,
+            prn_seed: int | None = None,
+            **init_kwargs
+    ) -> list[openmc.SourceParticle]:
+        """Sample external source and return source particles.
+
+        .. versionadded:: 0.15.1
+
+        Parameters
+        ----------
+        n_samples : int
+            Number of samples
+        prn_seed : int
+            Pseudorandom number generator (PRNG) seed; if None, one will be
+            generated randomly.
+        **init_kwargs
+            Keyword arguments passed to :func:`openmc.lib.init`
+
+        Returns
+        -------
+        list of openmc.SourceParticle
+            List of samples source particles
+        """
+        import openmc.lib
+
+        # Silence output by default. Also set arguments to start in volume
+        # calculation mode to avoid loading cross sections
+        init_kwargs.setdefault('output', False)
+        init_kwargs.setdefault('args', ['-c'])
+
+        with change_directory(tmpdir=True):
+            # Export model within temporary directory
+            self.export_to_model_xml()
+
+            # Sample external source sites
+            with openmc.lib.run_in_memory(**init_kwargs):
+                return openmc.lib.sample_external_source(
+                    n_samples=n_samples, prn_seed=prn_seed
+                )
+
     def plot_geometry(self, output=True, cwd='.', openmc_exec='openmc'):
         """Creates plot images as specified by the Model.plots attribute
 

openmc/universe.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import math
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
@@ -230,7 +231,7 @@ class Universe(UniverseBase):
         return self._cells
 
     @property
-    def bounding_box(self):
+    def bounding_box(self) -> openmc.BoundingBox:
         regions = [c.region for c in self.cells.values()
                    if c.region is not None]
         if regions:

tests/unit_tests/test_model.py

@@ -591,3 +591,32 @@ def test_single_xml_exec(run_in_tmpdir):
 
     os.mkdir('subdir')
     pincell_model.run(path='subdir')
+
+
+def test_model_plot():
+    # plots the geometry with source location and checks the resulting
+    # matplotlib includes the correct coordinates for the scatter plot for all
+    # basis.
+
+    surface = openmc.Sphere(r=600, boundary_type="vacuum")
+    cell = openmc.Cell(region=-surface)
+    geometry = openmc.Geometry([cell])
+    source = openmc.IndependentSource(space=openmc.stats.Point((1, 2, 3)))
+    settings = openmc.Settings(particles=1, batches=1, source=source)
+    model = openmc.Model(geometry, settings=settings)
+
+    plot = model.plot(n_samples=1, plane_tolerance=4.0, basis="xy")
+    coords = plot.axes.collections[0].get_offsets().data.flatten()
+    assert (coords == np.array([1.0, 2.0])).all()
+
+    plot = model.plot(n_samples=1, plane_tolerance=4.0, basis="xz")
+    coords = plot.axes.collections[0].get_offsets().data.flatten()
+    assert (coords == np.array([1.0, 3.0])).all()
+
+    plot = model.plot(n_samples=1, plane_tolerance=4.0, basis="yz")
+    coords = plot.axes.collections[0].get_offsets().data.flatten()
+    assert (coords == np.array([2.0, 3.0])).all()
+
+    plot = model.plot(n_samples=1, plane_tolerance=0.1, basis="xy")
+    coords = plot.axes.collections[0].get_offsets().data.flatten()
+    assert (coords == np.array([])).all()