from collections import Iterable
from numbers import Real, Integral
from xml.etree import ElementTree as ET
import sys
import warnings
import numpy as np
import openmc
import openmc.checkvalue as cv
from openmc.clean_xml import clean_xml_indentation
if sys.version_info[0] >= 3:
basestring = str
# A static variable for auto-generated Plot IDs
AUTO_PLOT_ID = 10000
def reset_auto_plot_id():
"""Reset counter for auto-generated plot IDs."""
global AUTO_PLOT_ID
AUTO_PLOT_ID = 10000
BASES = ['xy', 'xz', 'yz']
[docs]class Plot(object):
"""Definition of a finite region of space to be plotted, either as a slice plot
in two dimensions or as a voxel plot in three dimensions.
Parameters
----------
plot_id : int
Unique identifier for the plot
name : str
Name of the plot
Attributes
----------
id : int
Unique identifier
name : str
Name of the plot
width : Iterable of float
Width of the plot in each basis direction
pixels : Iterable of int
Number of pixels to use in each basis direction
origin : tuple or list of ndarray
Origin (center) of the plot
filename :
Path to write the plot to
color : {'cell', 'mat'}
Indicate whether the plot should be colored by cell or by material
type : {'slice', 'voxel'}
The type of the plot
basis : {'xy', 'xz', 'yz'}
The basis directions for the plot
background : tuple or list of ndarray
Color of the background defined by RGB
mask_components : Iterable of int
Unique id numbers of the cells or materials to plot
mask_background : Iterable of int
Color to apply to all cells/materials not listed in mask_components
defined by RGB
col_spec : dict
Dictionary indicating that certain cells/materials (keys) should be
colored with a specific RGB (values)
"""
def __init__(self, plot_id=None, name=''):
# Initialize Plot class attributes
self.id = plot_id
self.name = name
self._width = [4.0, 4.0]
self._pixels = [1000, 1000]
self._origin = [0., 0., 0.]
self._filename = 'plot'
self._color = 'cell'
self._type = 'slice'
self._basis = 'xy'
self._background = [0, 0, 0]
self._mask_components = None
self._mask_background = None
self._col_spec = None
@property
def id(self):
return self._id
@property
def name(self):
return self._name
@property
def width(self):
return self._width
@property
def pixels(self):
return self._pixels
@property
def origin(self):
return self._origin
@property
def filename(self):
return self._filename
@property
def color(self):
return self._color
@property
def type(self):
return self._type
@property
def basis(self):
return self._basis
@property
def background(self):
return self._background
@property
def mask_components(self):
return self._mask_components
@property
def mask_background(self):
return self._mask_background
@property
def col_spec(self):
return self._col_spec
@id.setter
def id(self, plot_id):
if plot_id is None:
global AUTO_PLOT_ID
self._id = AUTO_PLOT_ID
AUTO_PLOT_ID += 1
else:
cv.check_type('plot ID', plot_id, Integral)
cv.check_greater_than('plot ID', plot_id, 0, equality=True)
self._id = plot_id
@name.setter
def name(self, name):
cv.check_type('plot name', name, basestring)
self._name = name
@width.setter
def width(self, width):
cv.check_type('plot width', width, Iterable, Real)
cv.check_length('plot width', width, 2, 3)
self._width = width
@origin.setter
def origin(self, origin):
cv.check_type('plot origin', origin, Iterable, Real)
cv.check_length('plot origin', origin, 3)
self._origin = origin
@pixels.setter
def pixels(self, pixels):
cv.check_type('plot pixels', pixels, Iterable, Integral)
cv.check_length('plot pixels', pixels, 2, 3)
for dim in pixels:
cv.check_greater_than('plot pixels', dim, 0)
self._pixels = pixels
@filename.setter
def filename(self, filename):
cv.check_type('filename', filename, basestring)
self._filename = filename
@color.setter
def color(self, color):
cv.check_type('plot color', color, basestring)
cv.check_value('plot color', color, ['cell', 'mat'])
self._color = color
@type.setter
def type(self, plottype):
cv.check_type('plot type', plottype, basestring)
cv.check_value('plot type', plottype, ['slice', 'voxel'])
self._type = plottype
@basis.setter
def basis(self, basis):
cv.check_type('plot basis', basis, basestring)
cv.check_value('plot basis', basis, ['xy', 'xz', 'yz'])
self._basis = basis
@background.setter
def background(self, background):
cv.check_type('plot background', background, Iterable, Integral)
cv.check_length('plot background', background, 3)
for rgb in background:
cv.check_greater_than('plot background', rgb, 0, True)
cv.check_less_than('plot background', rgb, 256)
self._background = background
@col_spec.setter
def col_spec(self, col_spec):
cv.check_type('plot col_spec parameter', col_spec, dict, Integral)
for key in col_spec:
if key < 0:
msg = 'Unable to create Plot ID="{0}" with col_spec ID "{1}" ' \
'which is less than 0'.format(self._id, key)
raise ValueError(msg)
elif not isinstance(col_spec[key], Iterable):
msg = 'Unable to create Plot ID="{0}" with col_spec RGB values' \
' "{1}" which is not iterable'.format(self._id, col_spec[key])
raise ValueError(msg)
elif len(col_spec[key]) != 3:
msg = 'Unable to create Plot ID="{0}" with col_spec RGB ' \
'values of length "{1}" since 3 values must be ' \
'input'.format(self._id, len(col_spec[key]))
raise ValueError(msg)
self._col_spec = col_spec
@mask_components.setter
def mask_components(self, mask_components):
cv.check_type('plot mask_components', mask_components, Iterable, Integral)
for component in mask_components:
cv.check_greater_than('plot mask_components', component, 0, True)
self._mask_components = mask_components
@mask_background.setter
def mask_background(self, mask_background):
cv.check_type('plot mask background', mask_background, Iterable, Integral)
cv.check_length('plot mask background', mask_background, 3)
for rgb in mask_background:
cv.check_greater_than('plot mask background', rgb, 0, True)
cv.check_less_than('plot mask background', rgb, 256)
self._mask_background = mask_background
def __repr__(self):
string = 'Plot\n'
string += '{0: <16}{1}{2}\n'.format('\tID', '=\t', self._id)
string += '{0: <16}{1}{2}\n'.format('\tName', '=\t', self._name)
string += '{0: <16}{1}{2}\n'.format('\tFilename', '=\t', self._filename)
string += '{0: <16}{1}{2}\n'.format('\tType', '=\t', self._type)
string += '{0: <16}{1}{2}\n'.format('\tBasis', '=\t', self._basis)
string += '{0: <16}{1}{2}\n'.format('\tWidth', '=\t', self._width)
string += '{0: <16}{1}{2}\n'.format('\tOrigin', '=\t', self._origin)
string += '{0: <16}{1}{2}\n'.format('\tPixels', '=\t', self._origin)
string += '{0: <16}{1}{2}\n'.format('\tColor', '=\t', self._color)
string += '{0: <16}{1}{2}\n'.format('\tMask', '=\t',
self._mask_components)
string += '{0: <16}{1}{2}\n'.format('\tMask', '=\t',
self._mask_background)
string += '{0: <16}{1}{2}\n'.format('\tCol Spec', '=\t', self._col_spec)
return string
[docs] def colorize(self, geometry, seed=1):
"""Generate a color scheme for each domain in the plot.
This routine may be used to generate random, reproducible color schemes.
The colors generated are based upon cell/material IDs in the geometry.
Parameters
----------
geometry : openmc.Geometry
The geometry for which the plot is defined
seed : Integral
The random number seed used to generate the color scheme
"""
cv.check_type('geometry', geometry, openmc.Geometry)
cv.check_type('seed', seed, Integral)
cv.check_greater_than('seed', seed, 1, equality=True)
# Get collections of the domains which will be plotted
if self.color is 'mat':
domains = geometry.get_all_materials()
else:
domains = geometry.get_all_cells()
# Set the seed for the random number generator
np.random.seed(seed)
# Generate random colors for each feature
self.col_spec = {}
for domain in domains:
r = np.random.randint(0, 256)
g = np.random.randint(0, 256)
b = np.random.randint(0, 256)
self.col_spec[domain] = (r, g, b)
[docs] def highlight_domains(self, geometry, domains, seed=1,
alpha=0.5, background='gray'):
"""Use alpha compositing to highlight one or more domains in the plot.
This routine generates a color scheme and applies alpha compositing
to make all domains except the highlighted ones appear partially
transparent.
Parameters
----------
geometry : openmc.Geometry
The geometry for which the plot is defined
domains : Iterable of Integral
A collection of the domain IDs to highlight in the plot
seed : Integral
The random number seed used to generate the color scheme
alpha : Real in [0,1]
The value to apply in alpha compisiting
background : 3-tuple of Integral or 'white' or 'black' or 'gray'
The background color to apply in alpha compisiting
"""
cv.check_iterable_type('domains', domains, Integral)
cv.check_type('alpha', alpha, Real)
cv.check_greater_than('alpha', alpha, 0., equality=True)
cv.check_less_than('alpha', alpha, 1., equality=True)
# Get a background (R,G,B) tuple to apply in alpha compositing
if isinstance(background, basestring):
if background == 'white':
background = (255, 255, 255)
elif background == 'black':
background = (0, 0, 0)
elif background == 'gray':
background = (160, 160, 160)
else:
msg = 'The background "{}" is not defined'.format(background)
raise ValueError(msg)
cv.check_iterable_type('background', background, Integral)
# Generate a color scheme
self.colorize(geometry, seed)
# Apply alpha compositing to the colors for all domains
# other than those the user wishes to highlight
for domain_id in self.col_spec:
if domain_id not in domains:
r, g, b = self.col_spec[domain_id]
r = int(((1-alpha) * background[0]) + (alpha * r))
g = int(((1-alpha) * background[1]) + (alpha * g))
b = int(((1-alpha) * background[2]) + (alpha * b))
self._col_spec[domain_id] = (r, g, b)
[docs] def get_plot_xml(self):
"""Return XML representation of the plot
Returns
-------
element : xml.etree.ElementTree.Element
XML element containing plot data
"""
element = ET.Element("plot")
element.set("id", str(self._id))
element.set("filename", self._filename)
element.set("color", self._color)
element.set("type", self._type)
if self._type is 'slice':
element.set("basis", self._basis)
subelement = ET.SubElement(element, "origin")
subelement.text = ' '.join(map(str, self._origin))
subelement = ET.SubElement(element, "width")
subelement.text = ' '.join(map(str, self._width))
subelement = ET.SubElement(element, "pixels")
subelement.text = ' '.join(map(str, self._pixels))
if self._mask_background is not None:
subelement = ET.SubElement(element, "background")
subelement.text = ' '.join(map(str, self._background))
if self._col_spec is not None:
for key in self._col_spec:
subelement = ET.SubElement(element, "col_spec")
subelement.set("id", str(key))
subelement.set("rgb", ' '.join(map(
str, self._col_spec[key])))
if self._mask_components is not None:
subelement = ET.SubElement(element, "mask")
subelement.set("components", ' '.join(map(
str, self._mask_components)))
subelement.set("background", ' '.join(map(
str, self._mask_background)))
return element
[docs]class Plots(cv.CheckedList):
"""Collection of Plots used for an OpenMC simulation.
This class corresponds directly to the plots.xml input file. It can be
thought of as a normal Python list where each member is a :class:`Plot`. It
behaves like a list as the following example demonstrates:
>>> xz_plot = openmc.Plot()
>>> big_plot = openmc.Plot()
>>> small_plot = openmc.Plot()
>>> p = openmc.Plots((xz_plot, big_plot))
>>> p.append(small_plot)
>>> small_plot = p.pop()
Parameters
----------
plots : Iterable of openmc.Plot
Plots to add to the collection
"""
def __init__(self, plots=None):
super(Plots, self).__init__(Plot, 'plots collection')
self._plots_file = ET.Element("plots")
if plots is not None:
self += plots
[docs] def add_plot(self, plot):
"""Add a plot to the file.
.. deprecated:: 0.8
Use :meth:`Plots.append` instead.
Parameters
----------
plot : openmc.Plot
Plot to add
"""
warnings.warn("Plots.add_plot(...) has been deprecated and may be "
"removed in a future version. Use Plots.append(...) "
"instead.", DeprecationWarning)
self.append(plot)
[docs] def append(self, plot):
"""Append plot to collection
Parameters
----------
plot : openmc.Plot
Plot to append
"""
super(Plots, self).append(plot)
[docs] def insert(self, index, plot):
"""Insert plot before index
Parameters
----------
index : int
Index in list
plot : openmc.Plot
Plot to insert
"""
super(Plots, self).insert(index, plot)
[docs] def remove_plot(self, plot):
"""Remove a plot from the file.
.. deprecated:: 0.8
Use :meth:`Plots.remove` instead.
Parameters
----------
plot : openmc.Plot
Plot to remove
"""
warnings.warn("Plots.remove_plot(...) has been deprecated and may be "
"removed in a future version. Use Plots.remove(...) "
"instead.", DeprecationWarning)
self.remove(plot)
[docs] def colorize(self, geometry, seed=1):
"""Generate a consistent color scheme for each domain in each plot.
This routine may be used to generate random, reproducible color schemes.
The colors generated are based upon cell/material IDs in the geometry.
The color schemes will be consistent for all plots in "plots.xml".
Parameters
----------
geometry : openmc.Geometry
The geometry for which the plots are defined
seed : Integral
The random number seed used to generate the color scheme
"""
for plot in self:
plot.colorize(geometry, seed)
[docs] def highlight_domains(self, geometry, domains, seed=1,
alpha=0.5, background='gray'):
"""Use alpha compositing to highlight one or more domains in the plot.
This routine generates a color scheme and applies alpha compositing
to make all domains except the highlighted ones partially transparent.
Parameters
----------
geometry : openmc.Geometry
The geometry for which the plot is defined
domains : Iterable of Integral
A collection of the domain IDs to highlight in the plot
seed : Integral
The random number seed used to generate the color scheme
alpha : Real in [0,1]
The value to apply in alpha compisiting
background : 3-tuple of Integral or 'white' or 'black' or 'gray'
The background color to apply in alpha compisiting
"""
for plot in self:
plot.highlight_domains(geometry, domains, seed, alpha, background)
def _create_plot_subelements(self):
for plot in self:
xml_element = plot.get_plot_xml()
if len(plot.name) > 0:
self._plots_file.append(ET.Comment(plot.name))
self._plots_file.append(xml_element)
[docs] def export_to_xml(self):
"""Create a plots.xml file that can be used by OpenMC.
"""
# Reset xml element tree
self._plots_file.clear()
self._create_plot_subelements()
# Clean the indentation in the file to be user-readable
clean_xml_indentation(self._plots_file)
# Write the XML Tree to the plots.xml file
tree = ET.ElementTree(self._plots_file)
tree.write("plots.xml", xml_declaration=True,
encoding='utf-8', method="xml")