diff --git a/openmc/cell.py b/openmc/cell.py index 2e0f2ff13..cd0573e8b 100644 --- a/openmc/cell.py +++ b/openmc/cell.py @@ -10,6 +10,7 @@ import openmc import openmc.checkvalue as cv from ._xml import get_text from .mixin import IDManagerMixin +from .plots import add_plot_params from .region import Region, Complement from .surface import Halfspace from .bounding_box import BoundingBox @@ -559,64 +560,11 @@ class Cell(IDManagerMixin): return memo[self] + @add_plot_params def plot(self, *args, **kwargs): """Display a slice plot of the cell. .. versionadded:: 0.14.0 - - 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 - color_by : {'cell', 'material'} - Indicate whether the plot should be colored by cell or by material - colors : dict - Assigns colors to specific materials or cells. Keys are instances of - :class:`Cell` or :class:`Material` and values are RGB 3-tuples, RGBA - 4-tuples, or strings indicating SVG color names. Red, green, blue, - and alpha should all be floats in the range [0.0, 1.0], for example: - - .. code-block:: python - - # Make water blue - water = openmc.Cell(fill=h2o) - water.plot(colors={water: (0., 0., 1.)}) - seed : int - Seed for the random number generator - openmc_exec : str - Path to OpenMC executable. - axes : matplotlib.Axes - Axes to draw to - legend : bool - Whether a legend showing material or cell names should be drawn - legend_kwargs : dict - Keyword arguments passed to :func:`matplotlib.pyplot.legend`. - 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 - """ # Create dummy universe but preserve used_ids next_id = openmc.Universe.next_id diff --git a/openmc/geometry.py b/openmc/geometry.py index 28e1e48ec..c069d5796 100644 --- a/openmc/geometry.py +++ b/openmc/geometry.py @@ -9,6 +9,7 @@ import lxml.etree as ET import openmc import openmc._xml as xml +from .plots import add_plot_params from .checkvalue import check_type, check_less_than, check_greater_than, PathLike @@ -747,62 +748,10 @@ class Geometry: clone.root_universe = self.root_universe.clone() return clone + @add_plot_params def plot(self, *args, **kwargs): """Display a slice plot of the geometry. .. versionadded:: 0.14.0 - - 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 - color_by : {'cell', 'material'} - Indicate whether the plot should be colored by cell or by material - colors : dict - Assigns colors to specific materials or cells. Keys are instances of - :class:`Cell` or :class:`Material` and values are RGB 3-tuples, RGBA - 4-tuples, or strings indicating SVG color names. Red, green, blue, - and alpha should all be floats in the range [0.0, 1.0], for - example:: - - # Make water blue - water = openmc.Cell(fill=h2o) - universe.plot(..., colors={water: (0., 0., 1.)) - seed : int - Seed for the random number generator - openmc_exec : str - Path to OpenMC executable. - axes : matplotlib.Axes - Axes to draw to - legend : bool - Whether a legend showing material or cell names should be drawn - legend_kwargs : dict - Keyword arguments passed to :func:`matplotlib.pyplot.legend`. - 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 """ - return self.root_universe.plot(*args, **kwargs) diff --git a/openmc/model/model.py b/openmc/model/model.py index c8329c2fb..8dd13ef6b 100644 --- a/openmc/model/model.py +++ b/openmc/model/model.py @@ -1,9 +1,10 @@ from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from functools import lru_cache from pathlib import Path -from numbers import Integral -from tempfile import NamedTemporaryFile +import math +from numbers import Integral, Real +from tempfile import NamedTemporaryFile, TemporaryDirectory import warnings import h5py @@ -16,6 +17,7 @@ from openmc.dummy_comm import DummyCommunicator from openmc.executor import _process_CLI_arguments from openmc.checkvalue import check_type, check_value, PathLike from openmc.exceptions import InvalidIDError +from openmc.plots import add_plot_params from openmc.utility_funcs import change_directory @@ -130,6 +132,10 @@ class Model: for plot in plots: self._plots.append(plot) + @property + def bounding_box(self) -> openmc.BoundingBox: + return self.geometry.bounding_box + @property def is_initialized(self) -> bool: try: @@ -827,78 +833,207 @@ class Model: openmc.lib.materials[domain_id].volume = \ vol_calc.volumes[domain_id].n + @add_plot_params def plot( self, + origin: Sequence[float] | None = None, + width: Sequence[float] | None = None, + pixels: int | Sequence[int] = 40000, + basis: str = 'xy', + color_by: str = 'cell', + colors: dict | None = None, + seed: int | None = None, + openmc_exec: PathLike = 'openmc', + axes=None, + legend: bool = False, + axis_units: str = 'cm', + outline: bool | str = False, n_samples: int | None = None, plane_tolerance: float = 1., + legend_kwargs: dict | None = None, source_kwargs: dict | None = None, + contour_kwargs: dict | None = None, **kwargs, ): - """Display a slice plot of the geometry. + """Display a slice plot of the model. .. versionadded:: 0.15.1 - - Parameters - ---------- - n_samples : int, optional - 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, optional - 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 """ + import matplotlib.image as mpimg + import matplotlib.patches as mpatches + import matplotlib.pyplot as plt check_type('n_samples', n_samples, int | None) - check_type('plane_tolerance', plane_tolerance, float) + check_type('plane_tolerance', plane_tolerance, Real) + if legend_kwargs is None: + legend_kwargs = {} + legend_kwargs.setdefault('bbox_to_anchor', (1.05, 1)) + legend_kwargs.setdefault('loc', 2) + legend_kwargs.setdefault('borderaxespad', 0.0) if source_kwargs is None: source_kwargs = {} source_kwargs.setdefault('marker', 'x') - ax = self.geometry.plot(**kwargs) + # Determine extents of plot + if basis == 'xy': + x, y, z = 0, 1, 2 + xlabel, ylabel = f'x [{axis_units}]', f'y [{axis_units}]' + elif basis == 'yz': + x, y, z = 1, 2, 0 + xlabel, ylabel = f'y [{axis_units}]', f'z [{axis_units}]' + elif basis == 'xz': + x, y, z = 0, 2, 1 + xlabel, ylabel = f'x [{axis_units}]', f'z [{axis_units}]' + + bb = self.bounding_box + # checks to see if bounding box contains -inf or inf values + if np.isinf(bb.extent[basis]).any(): + if origin is None: + origin = (0, 0, 0) + if width is None: + width = (10, 10) + else: + if origin is None: + # if nan values in the bb.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(bb.center) + if width is None: + bb_width = bb.width + width = (bb_width[x], bb_width[y]) + + if isinstance(pixels, int): + aspect_ratio = width[0] / width[1] + pixels_y = math.sqrt(pixels / aspect_ratio) + pixels = (int(pixels / pixels_y), int(pixels_y)) + + axis_scaling_factor = {'km': 0.00001, 'm': 0.01, 'cm': 1, 'mm': 10} + + x_min = (origin[x] - 0.5*width[0]) * axis_scaling_factor[axis_units] + x_max = (origin[x] + 0.5*width[0]) * axis_scaling_factor[axis_units] + 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] + + with TemporaryDirectory() as tmpdir: + if seed is not None: + self.settings.plot_seed = seed + + # Create plot object matching passed arguments + plot = openmc.Plot() + plot.origin = origin + plot.width = width + plot.pixels = pixels + plot.basis = basis + plot.color_by = color_by + if colors is not None: + plot.colors = colors + self.plots.append(plot) + + # Run OpenMC in geometry plotting mode + self.plot_geometry(False, cwd=tmpdir, openmc_exec=openmc_exec) + + # Read image from file + img_path = Path(tmpdir) / f'plot_{plot.id}.png' + if not img_path.is_file(): + img_path = img_path.with_suffix('.ppm') + img = mpimg.imread(str(img_path)) + + # Create a figure sized such that the size of the axes within + # exactly matches the number of pixels specified + if axes is None: + px = 1/plt.rcParams['figure.dpi'] + fig, axes = plt.subplots() + axes.set_xlabel(xlabel) + axes.set_ylabel(ylabel) + params = fig.subplotpars + width = pixels[0]*px/(params.right - params.left) + height = pixels[1]*px/(params.top - params.bottom) + fig.set_size_inches(width, height) + + if outline: + # Combine R, G, B values into a single int + rgb = (img * 256).astype(int) + image_value = (rgb[..., 0] << 16) + \ + (rgb[..., 1] << 8) + (rgb[..., 2]) + + # Set default arguments for contour() + if contour_kwargs is None: + contour_kwargs = {} + contour_kwargs.setdefault('colors', 'k') + contour_kwargs.setdefault('linestyles', 'solid') + contour_kwargs.setdefault('algorithm', 'serial') + + axes.contour( + image_value, + origin="upper", + levels=np.unique(image_value), + extent=(x_min, x_max, y_min, y_max), + **contour_kwargs + ) + + # add legend showing which colors represent which material + # or cell if that was requested + if legend: + if plot.colors == {}: + raise ValueError("Must pass 'colors' dictionary if you " + "are adding a legend via legend=True.") + + if color_by == "cell": + expected_key_type = openmc.Cell + else: + expected_key_type = openmc.Material + + patches = [] + for key, color in plot.colors.items(): + + if isinstance(key, int): + raise TypeError( + "Cannot use IDs in colors dict for auto legend.") + elif not isinstance(key, expected_key_type): + raise TypeError( + "Color dict key type does not match color_by") + + # this works whether we're doing cells or materials + label = key.name if key.name != '' else key.id + + # matplotlib takes RGB on 0-1 scale rather than 0-255. at + # this point PlotBase has already checked that 3-tuple + # based colors are already valid, so if the length is three + # then we know it just needs to be converted to the 0-1 + # format. + if len(color) == 3 and not isinstance(color, str): + scaled_color = ( + color[0]/255, color[1]/255, color[2]/255) + else: + scaled_color = color + + key_patch = mpatches.Patch(color=scaled_color, label=label) + patches.append(key_patch) + + axes.legend(handles=patches, **legend_kwargs) + + # Plot image and return the axes + if outline != 'only': + axes.imshow(img, extent=(x_min, x_max, y_min, y_max), **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] - + # Get points within tolerance of the slice plane + slice_value = origin[z] 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 + if (slice_value - tol < particle.r[z] < slice_value + tol): + xs.append(particle.r[x]) + ys.append(particle.r[y]) + axes.scatter(xs, ys, **source_kwargs) + + return axes def sample_external_source( self, diff --git a/openmc/plots.py b/openmc/plots.py index 2bb3ef65f..a59a7f0eb 100644 --- a/openmc/plots.py +++ b/openmc/plots.py @@ -165,6 +165,91 @@ _SVG_COLORS = { 'yellowgreen': (154, 205, 50) } +_PLOT_PARAMS = """ + Parameters + ---------- + origin : iterable of float + Coordinates at the origin of the plot. If left as None, + the center of the bounding box will be used to attempt to ascertain + the origin with infinite values being replaced by 0. + width : iterable of float + Width of the plot in each basis direction. If left as none then the + width of the bounding box 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 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 + color_by : {'cell', 'material'} + Indicate whether the plot should be colored by cell or by material + colors : dict + Assigns colors to specific materials or cells. Keys are instances of + :class:`Cell` or :class:`Material` and values are RGB 3-tuples, RGBA + 4-tuples, or strings indicating SVG color names. Red, green, blue, + and alpha should all be floats in the range [0.0, 1.0], for example: + + .. code-block:: python + + # Make water blue + water = openmc.Cell(fill=h2o) + universe.plot(..., colors={water: (0., 0., 1.)) + seed : int + Seed for the random number generator + openmc_exec : str + Path to OpenMC executable. + axes : matplotlib.Axes + Axes to draw to + + .. versionadded:: 0.13.1 + legend : bool + Whether a legend showing material or cell names should be drawn + + .. versionadded:: 0.14.0 + outline : bool or str + Whether outlines between color boundaries should be drawn. If set to + 'only', only outlines will be drawn. + + .. versionadded:: 0.14.0 + axis_units : {'km', 'm', 'cm', 'mm'} + Units used on the plot axis + + .. versionadded:: 0.14.0 + n_samples : int, optional + 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 + legend_kwargs : dict + Keyword arguments passed to :func:`matplotlib.pyplot.legend`. + + .. versionadded:: 0.14.0 + source_kwargs : dict, optional + Keyword arguments passed to :func:`matplotlib.pyplot.scatter`. + contour_kwargs : dict, optional + Keyword arguments passed to :func:`matplotlib.pyplot.contour`. + **kwargs + Keyword arguments passed to :func:`matplotlib.pyplot.imshow`. + + Returns + ------- + matplotlib.axes.Axes + Axes containing resulting image +""" + + +# Decorator for consistently adding plot parameters to docstrings (Model.plot, +# Geometry.plot, Universe.plot, etc.) +def add_plot_params(func): + func.__doc__ += _PLOT_PARAMS + return func + def _get_plot_image(plot, cwd): from IPython.display import Image diff --git a/openmc/region.py b/openmc/region.py index e1cb83475..cb9f3abd2 100644 --- a/openmc/region.py +++ b/openmc/region.py @@ -8,6 +8,7 @@ import numpy as np import openmc from .bounding_box import BoundingBox +from .plots import add_plot_params class Region(ABC): @@ -343,47 +344,11 @@ class Region(ABC): return type(self)(n.rotate(rotation, pivot=pivot, order=order, inplace=inplace, memo=memo) for n in self) + @add_plot_params def plot(self, *args, **kwargs): """Display a slice plot of the region. .. versionadded:: 0.15.0 - - 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: diff --git a/openmc/universe.py b/openmc/universe.py index 71ffa726b..02b794dee 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -1,17 +1,14 @@ from __future__ import annotations -import math from abc import ABC, abstractmethod from collections.abc import Iterable from numbers import Real -from pathlib import Path -from tempfile import TemporaryDirectory -import warnings import numpy as np import openmc import openmc.checkvalue as cv from .mixin import IDManagerMixin +from .plots import add_plot_params class UniverseBase(ABC, IDManagerMixin): @@ -334,234 +331,21 @@ class UniverseBase(ABC, IDManagerMixin): return [self, cell] + cell.fill.find(p) return [] - # default kwargs that are passed to plt.legend in the plot method below. - _default_legend_kwargs = {'bbox_to_anchor': ( - 1.05, 1), 'loc': 2, 'borderaxespad': 0.0} - - def plot(self, origin=None, width=None, pixels=40000, - basis='xy', color_by='cell', colors=None, seed=None, - openmc_exec='openmc', axes=None, legend=False, axis_units='cm', - legend_kwargs=_default_legend_kwargs, outline=False, - **kwargs): + @add_plot_params + def plot(self, *args, **kwargs): """Display a slice plot of the universe. - - Parameters - ---------- - origin : iterable of float - Coordinates at the origin of the plot. If left as None, - universe.bounding_box.center will be used to attempt to ascertain - the origin with infinite values being replaced by 0. - 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 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 - color_by : {'cell', 'material'} - Indicate whether the plot should be colored by cell or by material - colors : dict - Assigns colors to specific materials or cells. Keys are instances of - :class:`Cell` or :class:`Material` and values are RGB 3-tuples, RGBA - 4-tuples, or strings indicating SVG color names. Red, green, blue, - and alpha should all be floats in the range [0.0, 1.0], for example: - - .. code-block:: python - - # Make water blue - water = openmc.Cell(fill=h2o) - universe.plot(..., colors={water: (0., 0., 1.)) - seed : int - Seed for the random number generator - openmc_exec : str - Path to OpenMC executable. - axes : matplotlib.Axes - Axes to draw to - - .. versionadded:: 0.13.1 - legend : bool - Whether a legend showing material or cell names should be drawn - - .. versionadded:: 0.14.0 - legend_kwargs : dict - Keyword arguments passed to :func:`matplotlib.pyplot.legend`. - - .. versionadded:: 0.14.0 - outline : bool - Whether outlines between color boundaries should be drawn - - .. versionadded:: 0.14.0 - axis_units : {'km', 'm', 'cm', 'mm'} - Units used on the plot axis - - .. versionadded:: 0.14.0 - **kwargs - Keyword arguments passed to :func:`matplotlib.pyplot.imshow` - - Returns - ------- - matplotlib.axes.Axes - Axes containing resulting image - """ - import matplotlib.image as mpimg - import matplotlib.patches as mpatches - import matplotlib.pyplot as plt + model = openmc.Model() + model.geometry = openmc.Geometry(self) - # Determine extents of plot - if basis == 'xy': - x, y = 0, 1 - xlabel, ylabel = f'x [{axis_units}]', f'y [{axis_units}]' - elif basis == 'yz': - x, y = 1, 2 - xlabel, ylabel = f'y [{axis_units}]', f'z [{axis_units}]' - elif basis == 'xz': - x, y = 0, 2 - xlabel, ylabel = f'x [{axis_units}]', f'z [{axis_units}]' + # 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 - bb = self.bounding_box - # checks to see if bounding box contains -inf or inf values - if np.isinf(bb.extent[basis]).any(): - if origin is None: - origin = (0, 0, 0) - if width is None: - width = (10, 10) - else: - if origin is None: - # if nan values in the bb.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(bb.center) - if width is None: - bb_width = bb.width - x_width = bb_width['xyz'.index(basis[0])] - y_width = bb_width['xyz'.index(basis[1])] - width = (x_width, y_width) - - if isinstance(pixels, int): - aspect_ratio = width[0] / width[1] - pixels_y = math.sqrt(pixels / aspect_ratio) - pixels = (int(pixels / pixels_y), int(pixels_y)) - - axis_scaling_factor = {'km': 0.00001, 'm': 0.01, 'cm': 1, 'mm': 10} - - x_min = (origin[x] - 0.5*width[0]) * axis_scaling_factor[axis_units] - x_max = (origin[x] + 0.5*width[0]) * axis_scaling_factor[axis_units] - 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] - - with TemporaryDirectory() as tmpdir: - model = openmc.Model() - model.geometry = openmc.Geometry(self) - if seed is not None: - model.settings.plot_seed = seed - - # 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 - - # Create plot object matching passed arguments - plot = openmc.Plot() - plot.origin = origin - plot.width = width - plot.pixels = pixels - plot.basis = basis - plot.color_by = color_by - if colors is not None: - plot.colors = colors - model.plots.append(plot) - - # Run OpenMC in geometry plotting mode - model.plot_geometry(False, cwd=tmpdir, openmc_exec=openmc_exec) - - # Read image from file - img_path = Path(tmpdir) / f'plot_{plot.id}.png' - if not img_path.is_file(): - img_path = img_path.with_suffix('.ppm') - img = mpimg.imread(str(img_path)) - - # Create a figure sized such that the size of the axes within - # exactly matches the number of pixels specified - if axes is None: - px = 1/plt.rcParams['figure.dpi'] - fig, axes = plt.subplots() - axes.set_xlabel(xlabel) - axes.set_ylabel(ylabel) - params = fig.subplotpars - width = pixels[0]*px/(params.right - params.left) - height = pixels[1]*px/(params.top - params.bottom) - fig.set_size_inches(width, height) - - if outline: - # Combine R, G, B values into a single int - rgb = (img * 256).astype(int) - image_value = (rgb[..., 0] << 16) + \ - (rgb[..., 1] << 8) + (rgb[..., 2]) - - axes.contour( - image_value, - origin="upper", - colors="k", - linestyles="solid", - linewidths=1, - levels=np.unique(image_value), - extent=(x_min, x_max, y_min, y_max), - ) - - # add legend showing which colors represent which material - # or cell if that was requested - if legend: - if plot.colors == {}: - raise ValueError("Must pass 'colors' dictionary if you " - "are adding a legend via legend=True.") - - if color_by == "cell": - expected_key_type = openmc.Cell - else: - expected_key_type = openmc.Material - - patches = [] - for key, color in plot.colors.items(): - - if isinstance(key, int): - raise TypeError( - "Cannot use IDs in colors dict for auto legend.") - elif not isinstance(key, expected_key_type): - raise TypeError( - "Color dict key type does not match color_by") - - # this works whether we're doing cells or materials - label = key.name if key.name != '' else key.id - - # matplotlib takes RGB on 0-1 scale rather than 0-255. at - # this point PlotBase has already checked that 3-tuple - # based colors are already valid, so if the length is three - # then we know it just needs to be converted to the 0-1 - # format. - if len(color) == 3 and not isinstance(color, str): - scaled_color = ( - color[0]/255, color[1]/255, color[2]/255) - else: - scaled_color = color - - key_patch = mpatches.Patch(color=scaled_color, label=label) - patches.append(key_patch) - - axes.legend(handles=patches, **legend_kwargs) - - # Plot image and return the axes - axes.imshow(img, extent=(x_min, x_max, y_min, y_max), **kwargs) - return axes + return model.plot(*args, **kwargs) def get_nuclides(self): """Returns all nuclides in the universe