From 8b77a8dd3b3a1c65b1e08f7f032bfee7c86ee51d Mon Sep 17 00:00:00 2001 From: Jonathan Shimwell Date: Fri, 27 Sep 2024 00:20:40 +0200 Subject: [PATCH] Adding source option to plot (#2863) Co-authored-by: Jon Shimwell Co-authored-by: Paul Romano --- openmc/geometry.py | 5 +- openmc/lib/core.py | 9 ++- openmc/model/model.py | 116 +++++++++++++++++++++++++++++++++ openmc/universe.py | 3 +- tests/unit_tests/test_model.py | 29 +++++++++ 5 files changed, 155 insertions(+), 7 deletions(-) diff --git a/openmc/geometry.py b/openmc/geometry.py index 6cce4c18c..a52dc4418 100644 --- a/openmc/geometry.py +++ b/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 diff --git a/openmc/lib/core.py b/openmc/lib/core.py index a9a549fa0..e646a9ae1 100644 --- a/openmc/lib/core.py +++ b/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: diff --git a/openmc/model/model.py b/openmc/model/model.py index 2ea579ab7..78f743a03 100644 --- a/openmc/model/model.py +++ b/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 diff --git a/openmc/universe.py b/openmc/universe.py index 9fab9ae51..d424c243b 100644 --- a/openmc/universe.py +++ b/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: diff --git a/tests/unit_tests/test_model.py b/tests/unit_tests/test_model.py index 404235bef..4b567c56d 100644 --- a/tests/unit_tests/test_model.py +++ b/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()