Merge pull request #2019 from paulromano/universe-plot-backend

Use `openmc` executable when calling Universe.plot
This commit is contained in:
Patrick Shriwise 2022-04-07 17:21:47 -05:00 committed by GitHub
commit 0c0e70fa39
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 64 additions and 92 deletions

View file

@ -442,7 +442,7 @@ class Decay(EqualityMixin):
# Read continuous spectrum
ci = {}
params, ci['probability'] = get_tab1_record(file_obj)
ci['type'] = get_decay_modes(params[0])
ci['from_mode'] = get_decay_modes(params[0])
# Read covariance (Ek, Fk) table
LCOV = params[3]

View file

@ -166,13 +166,13 @@ def plot_inline(plots, openmc_exec='openmc', cwd='.'):
plots = [plots]
# Create plots.xml
openmc.Plots(plots).export_to_xml()
openmc.Plots(plots).export_to_xml(cwd)
# Run OpenMC in geometry plotting mode
plot_geometry(False, openmc_exec, cwd)
if plots is not None:
images = [_get_plot_image(p) for p in plots]
images = [_get_plot_image(p, cwd) for p in plots]
display(*images)

View file

@ -395,7 +395,7 @@ class Model:
depletion_operator.cleanup_when_done = True
depletion_operator.finalize()
def export_to_xml(self, directory='.'):
def export_to_xml(self, directory='.', remove_surfs=False):
"""Export model to XML files.
Parameters
@ -403,7 +403,11 @@ class Model:
directory : str
Directory to write XML files to. If it doesn't exist already, it
will be created.
remove_surfs : bool
Whether or not to remove redundant surfaces from the geometry when
exporting.
.. versionadded:: 0.13.1
"""
# Create directory if required
d = Path(directory)
@ -411,7 +415,7 @@ class Model:
d.mkdir(parents=True)
self.settings.export_to_xml(d)
self.geometry.export_to_xml(d)
self.geometry.export_to_xml(d, remove_surfs=remove_surfs)
# If a materials collection was specified, export it. Otherwise, look
# for all materials in the geometry and use that to automatically build

View file

@ -164,16 +164,16 @@ _SVG_COLORS = {
}
def _get_plot_image(plot):
def _get_plot_image(plot, cwd):
from IPython.display import Image
# Make sure .png file was created
stem = plot.filename if plot.filename is not None else f'plot_{plot.id}'
png_file = f'{stem}.png'
if not Path(png_file).exists():
png_file = Path(cwd) / f'{stem}.png'
if not png_file.exists():
raise FileNotFoundError(f"Could not find .png image for plot {plot.id}")
return Image(png_file)
return Image(str(png_file))
class Plot(IDManagerMixin):
@ -784,13 +784,13 @@ class Plot(IDManagerMixin):
"""
# Create plots.xml
Plots([self]).export_to_xml()
Plots([self]).export_to_xml(cwd)
# Run OpenMC in geometry plotting mode
openmc.plot_geometry(False, openmc_exec, cwd)
# Return produced image
return _get_plot_image(self)
return _get_plot_image(self, cwd)
class Plots(cv.CheckedList):

View file

@ -1,9 +1,10 @@
from abc import ABC, abstractmethod
from collections import OrderedDict
from collections.abc import Iterable
from copy import copy, deepcopy
from copy import deepcopy
from numbers import Real
import random
from pathlib import Path
from tempfile import TemporaryDirectory
from xml.etree import ElementTree as ET
import numpy as np
@ -267,13 +268,9 @@ class Universe(UniverseBase):
def plot(self, origin=(0., 0., 0.), width=(1., 1.), pixels=(200, 200),
basis='xy', color_by='cell', colors=None, seed=None,
**kwargs):
openmc_exec='openmc', **kwargs):
"""Display a slice plot of the universe.
To display or save the plot, call :func:`matplotlib.pyplot.show` or
:func:`matplotlib.pyplot.savefig`. In a Jupyter notebook, enabling the
matplotlib inline backend will show the plot inline.
Parameters
----------
origin : Iterable of float
@ -297,14 +294,14 @@ class Universe(UniverseBase):
# 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.
seed : hashable object or None
Hashable object which is used to seed the random number generator
used to select colors. If None, the generator is seeded from the
current time.
.. versionadded:: 0.13.1
**kwargs
All keyword arguments are passed to
:func:`matplotlib.pyplot.imshow`.
Keyword arguments passed to :func:`matplotlib.pyplot.imshow`
Returns
-------
@ -312,83 +309,54 @@ class Universe(UniverseBase):
Resulting image
"""
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
# Seed the random number generator
if seed is not None:
random.seed(seed)
if colors is None:
# Create default dictionary if none supplied
colors = {}
else:
# Convert to RGBA if necessary
colors = copy(colors)
for obj, color in colors.items():
if isinstance(color, str):
if color.lower() not in _SVG_COLORS:
raise ValueError(f"'{color}' is not a valid color.")
colors[obj] = [x/255 for x in
_SVG_COLORS[color.lower()]] + [1.0]
elif len(color) == 3:
colors[obj] = list(color) + [1.0]
# Determine extents of plot
if basis == 'xy':
x_min = origin[0] - 0.5*width[0]
x_max = origin[0] + 0.5*width[0]
y_min = origin[1] - 0.5*width[1]
y_max = origin[1] + 0.5*width[1]
x, y = 0, 1
elif basis == 'yz':
# The x-axis will correspond to physical y and the y-axis will
# correspond to physical z
x_min = origin[1] - 0.5*width[0]
x_max = origin[1] + 0.5*width[0]
y_min = origin[2] - 0.5*width[1]
y_max = origin[2] + 0.5*width[1]
x, y = 1, 2
elif basis == 'xz':
# The y-axis will correspond to physical z
x_min = origin[0] - 0.5*width[0]
x_max = origin[0] + 0.5*width[0]
y_min = origin[2] - 0.5*width[1]
y_max = origin[2] + 0.5*width[1]
x, y = 0, 2
x_min = origin[x] - 0.5*width[0]
x_max = origin[x] + 0.5*width[0]
y_min = origin[y] - 0.5*width[1]
y_max = origin[y] + 0.5*width[1]
# Determine locations to determine cells at
x_coords = np.linspace(x_min, x_max, pixels[0], endpoint=False) + \
0.5*(x_max - x_min)/pixels[0]
y_coords = np.linspace(y_max, y_min, pixels[1], endpoint=False) - \
0.5*(y_max - y_min)/pixels[1]
with TemporaryDirectory() as tmpdir:
model = openmc.Model()
model.geometry = openmc.Geometry(self)
if seed is not None:
model.settings.seed = seed
# Initialize output image in RGBA format. Flip the pixels from
# traditional (x, y) to (y, x) used in graphics.
img = np.zeros((pixels[1], pixels[0], 4))
for i, x in enumerate(x_coords):
for j, y in enumerate(y_coords):
if basis == 'xy':
path = self.find((x, y, origin[2]))
elif basis == 'yz':
path = self.find((origin[0], x, y))
elif basis == 'xz':
path = self.find((x, origin[1], y))
# 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
plot.colors = colors
model.plots.append(plot)
if len(path) > 0:
try:
if color_by == 'cell':
obj = path[-1]
elif color_by == 'material':
if path[-1].fill_type == 'material':
obj = path[-1].fill
else:
continue
except AttributeError:
continue
if obj not in colors:
colors[obj] = (random.random(), random.random(),
random.random(), 1.0)
img[j, i, :] = colors[obj]
# Run OpenMC in geometry plotting mode
model.plot_geometry(False, cwd=tmpdir, openmc_exec=openmc_exec)
# Display image
return plt.imshow(img, extent=(x_min, x_max, y_min, y_max),
interpolation='nearest', **kwargs)
# Read image from file
img = mpimg.imread(Path(tmpdir) / f'plot_{plot.id}.png')
# Create a figure sized such that the size of the axes within
# exactly matches the number of pixels specified
px = 1/plt.rcParams['figure.dpi']
fig, ax = plt.subplots()
params = fig.subplotpars
width = pixels[0]*px/(params.right - params.left)
height = pixels[0]*px/(params.top - params.bottom)
fig.set_size_inches(width, height)
# Plot image and return the axes
return ax.imshow(img, extent=(x_min, x_max, y_min, y_max), **kwargs)
def add_cell(self, cell):
"""Add a cell to the universe.