crp/OpenMC
1
0
Fork 0
mirror of https://github.com/openmc-dev/openmc.git synced 2026-07-27 21:55:41 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

Provide a way to get ID maps from plot parameters on the Model class (#3481)

Browse source 
Co-authored-by: Jonathan Shimwell <drshimwell@gmail.com>
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
This commit is contained in:
Patrick Shriwise 2025-07-16 08:31:49 -05:00 • committed by GitHub
parent 58ee8d825d
commit 6372c29cfa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 378 additions and 52 deletions
Download patch file Download diff file Expand all files Collapse all files

1
openmc/lib/plot.py
View file

@ -90,6 +90,7 @@ class _PlotBase(Structure):
def __init__(self):
self.level_ = -1
self.basis_ = 1
self.color_overlaps_ = False
@property

173
openmc/model/model.py
View file

@ -20,7 +20,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.plots import add_plot_params, _BASIS_INDICES
from openmc.utility_funcs import change_directory
@ -902,6 +902,111 @@ class Model:
openmc.lib.materials[domain_id].volume = \
vol_calc.volumes[domain_id].n
def _set_plot_defaults(
self,
origin: Sequence[float] | None,
width: Sequence[float] | None,
pixels: int | Sequence[int],
basis: str
):
x, y, _ = _BASIS_INDICES[basis]
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))
return origin, width, pixels
def id_map(
self,
origin: Sequence[float] | None = None,
width: Sequence[float] | None = None,
pixels: int | Sequence[int] = 40000,
basis: str = 'xy',
**init_kwargs
) -> np.ndarray:
"""Generate an ID map for domains based on the plot parameters
If the model is not yet initialized, it will be initialized with
openmc.lib. If the model is initialized, the model will remain
initialized after this method call exits.
.. versionadded:: 0.15.3
Parameters
----------
origin : Sequence[float], optional
Origin of the plot. If unspecified, this argument defaults to the
center of the bounding box if the bounding box does not contain inf
values for the provided basis, otherwise (0.0, 0.0, 0.0).
width : Sequence[float], optional
Width of the plot. If unspecified, this argument defaults to the
width of the bounding box if the bounding box does not contain inf
values for the provided basis, otherwise (10.0, 10.0).
pixels : int | Sequence[int], optional
If an iterable of ints is provided then this directly sets the
number of pixels to use in each basis direction. If a single int is
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 based on the width argument.
basis : {'xy', 'yz', 'xz'}, optional
Basis of the plot.
**init_kwargs
Keyword arguments passed to :meth:`Model.init_lib`.
Returns
-------
id_map : numpy.ndarray
A NumPy array with shape (vertical pixels, horizontal pixels, 3) of
OpenMC property IDs with dtype int32. The last dimension of the
array contains cell IDs, cell instances, and material IDs (in that
order).
"""
import openmc.lib
origin, width, pixels = self._set_plot_defaults(
origin, width, pixels, basis)
# initialize the openmc.lib.plot._PlotBase object
plot_obj = openmc.lib.plot._PlotBase()
plot_obj.origin = origin
plot_obj.width = width[0]
plot_obj.height = width[1]
plot_obj.h_res = pixels[0]
plot_obj.v_res = pixels[1]
plot_obj.basis = basis
if self.is_initialized:
return openmc.lib.id_map(plot_obj)
else:
# 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 openmc.lib.TemporarySession(self, **init_kwargs):
return openmc.lib.id_map(plot_obj)
@add_plot_params
def plot(
self,
@ -945,39 +1050,13 @@ class Model:
source_kwargs = {}
source_kwargs.setdefault('marker', 'x')
# Set indices using basis and create axis labels
x, y, z = _BASIS_INDICES[basis]
xlabel, ylabel = f'{basis[0]} [{axis_units}]', f'{basis[1]} [{axis_units}]'
# 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))
origin, width, pixels = self._set_plot_defaults(
origin, width, pixels, basis)
axis_scaling_factor = {'km': 0.00001, 'm': 0.01, 'cm': 1, 'mm': 10}
@ -1124,10 +1203,10 @@ class Model:
return axes
def sample_external_source(
self,
n_samples: int = 1000,
prn_seed: int | None = None,
**init_kwargs
self,
n_samples: int = 1000,
prn_seed: int | None = None,
**init_kwargs
) -> openmc.ParticleList:
"""Sample external source and return source particles.
@ -1150,17 +1229,17 @@ class Model:
"""
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'])
if self.is_initialized:
return openmc.lib.sample_external_source(
n_samples=n_samples, prn_seed=prn_seed
)
else:
# 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):
with openmc.lib.TemporarySession(self, **init_kwargs):
return openmc.lib.sample_external_source(
n_samples=n_samples, prn_seed=prn_seed
)

13
openmc/plots.py
View file

@ -15,6 +15,8 @@ from .mixin import IDManagerMixin
_BASES = {'xy', 'xz', 'yz'}
_BASIS_INDICES = {'xy': (0, 1, 2), 'xz': (0, 2, 1), 'yz': (1, 2, 0)}
_SVG_COLORS = {
'aliceblue': (240, 248, 255),
'antiquewhite': (250, 235, 215),
@ -178,11 +180,12 @@ _PLOT_PARAMS = """
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.
If an iterable of ints is provided then this directly sets the
number of pixels to use in each basis direction. If a single int
is 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 based on the width
argument.
basis : {'xy', 'xz', 'yz'}
The basis directions for the plot
color_by : {'cell', 'material'}

243
tests/unit_tests/test_model.py
View file

@ -649,3 +649,246 @@ def test_model_plot():
# ensure that all of the data in the image data is either white or red
test_mask = (image_data == white) | (image_data == red)
assert np.all(test_mask), "Colors other than white or red found in overlap plot image"
def test_model_id_map_initialization(run_in_tmpdir):
model = openmc.examples.pwr_assembly()
model.init_lib(output=False)
id_map = model.id_map(
pixels=(100, 100),
basis='xy',
origin=(0, 0, 0),
width=(10, 10),
)
assert id_map.shape == (100, 100, 3)
assert id_map.dtype == np.int32
max_cell_id = max(model.geometry.get_all_cells().keys())
max_material_id = max(model.geometry.get_all_materials().keys())
# add some spot checks for the id_map
# Check that the array contains valid cell/material IDs (not all -2)
# The -2 values indicate outside the geometry
assert not np.all(id_map == -2), "All values are -2, indicating no valid geometry found"
# Check that we have valid cell IDs (first dimension)
valid_cell_ids = id_map[:, :, 0]
assert np.any(valid_cell_ids >= 0), "No valid cell IDs found in the id_map"
# Check that we have valid material IDs (third dimension)
valid_material_ids = id_map[:, :, 2]
assert np.any(valid_material_ids >= 0), "No valid material IDs found in the id_map"
# Check that the middle dimension (cell instances) is consistent
# Cell instances should be >= 0 when cell IDs are valid
cell_instances = id_map[:, :, 1]
valid_cells = valid_cell_ids >= 0
if np.any(valid_cells):
assert np.all(cell_instances[valid_cells] >= 0), "Invalid cell instances found for valid cells"
# Check that the array contains reasonable ranges of values
# Cell IDs should be within the expected range for the assembly
if np.any(valid_cell_ids >= 0):
max_map_cell_id = np.max(valid_cell_ids)
assert max_map_cell_id <= max_cell_id, \
f"Cell ID {max_map_cell_id} in the map is greater than the maximum cell ID {max_cell_id}"
# Material IDs should be within the expected range
if np.any(valid_material_ids >= 0):
max_map_material_id = np.max(valid_material_ids)
assert max_map_material_id <= max_material_id, \
f"Material ID {max_map_material_id} in the map is greater than the maximum material ID {max_material_id}"
# Test id_map with pixels outside the model geometry
# Use a plot that's far from the model center to ensure we get -2 values
outside_id_map = model.id_map(
pixels=(50, 50),
basis='xy',
origin=(1000, 1000, 0), # Far from the model center
width=(10, 10),
)
assert outside_id_map.shape == (50, 50, 3)
assert outside_id_map.dtype == np.int32
# All values should be -2 (outside geometry) for this plot
assert np.all(outside_id_map == -2), "Expected all values to be -2 for plot outside model geometry"
# Verify that the outside plot has the correct structure
assert np.all(outside_id_map[:, :, 0] == -2), "Cell IDs should all be -2 outside geometry"
assert np.all(outside_id_map[:, :, 1] == -2), "Cell instances should all be -2 outside geometry"
assert np.all(outside_id_map[:, :, 2] == -2), "Material IDs should all be -2 outside geometry"
# if the model is already initialized, it should not be finalized
# after calling this method
model.id_map(
pixels=(100, 100),
basis='xy',
origin=(0, 0, 0),
width=(10, 10),
)
assert model.is_initialized
# if the model is not initialized, it should be finalized
# before exiting this method
model.finalize_lib()
model.id_map(
pixels=(100, 100),
basis='xy',
origin=(0, 0, 0),
width=(10, 10),
)
assert not model.is_initialized
def test_id_map_aligned_model():
"""Test id_map with a 2x2 lattice where pixel boundaries align to cell boundaries"""
# Create materials -- identical compositions, different IDs
mat1 = openmc.Material(material_id=1, name='Material 1')
mat1.set_density('g/cm3', 1.0)
mat1.add_element('H', 1.0)
mat2 = openmc.Material(material_id=2, name='Material 2')
mat2.set_density('g/cm3', 1.0)
mat2.add_element('H', 1.0)
mat3 = openmc.Material(material_id=3, name='Material 3')
mat3.set_density('g/cm3', 1.0)
mat3.add_element('H', 1.0)
mat4 = openmc.Material(material_id=4, name='Material 4')
mat4.set_density('g/cm3', 1.0)
mat4.add_element('H', 1.0)
outer_mat = openmc.Material(material_id=5, name='Material 5')
outer_mat.set_density('g/cm3', 1.0)
outer_mat.add_element('H', 1.0)
inner_materials = [mat1, mat2, mat3, mat4]
# Create square surface that fits inside the lattice cell
# Lattice cell is 1 cm x 1 cm, so square will be 0.6 cm x 0.6 cm centered on the origin
square = openmc.model.RectangularPrism(0.6, 0.6, boundary_type='transmission')
# Create cells for this universe
inner_cell = openmc.Cell(cell_id=10, region=-square, name='inner_cell')
inner_cell.fill = inner_materials
outer_cell = openmc.Cell(cell_id=20, region=+square, name='outer_cell')
outer_cell.fill = outer_mat
# Create universe
universe = openmc.Universe(universe_id=100, cells=[inner_cell, outer_cell])
# Create 2x2 lattice
lattice = openmc.RectLattice(lattice_id=1)
lattice.lower_left = [-1.0, -1.0]
lattice.pitch = [1.0, 1.0]
lattice.universes = [[universe, universe], [universe, universe]]
# Create outer boundary
outer_boundary = openmc.model.RectangularPrism(2.0, 2.0, boundary_type='vacuum')
# Create root cell
root_cell = openmc.Cell(cell_id=1, name='root', fill=lattice, region=-outer_boundary)
# Create geometry
geometry = openmc.Geometry([root_cell])
# Create settings
settings = openmc.Settings()
settings.particles = 1000
settings.batches = 10
# Create model
model = openmc.Model(settings=settings, geometry=geometry)
# Generate id_map with pixel boundaries aligned to cell boundaries
# The model is 2 cm x 2 cm, so we'll use 200x200 pixels to get 0.01 cm resolution
# This allows us to align pixels with the squares inside each lattice cell
id_map = model.id_map(
pixels=(200, 200),
basis='xy',
origin=(0.0, 0.0, 0.0), # Align with lattice lower_left
width=(2.0, 2.0), # Align with lattice size
)
# Verify id_map properties
assert id_map.shape == (200, 200, 3)
assert id_map.dtype == np.int32
cell_id_map = id_map[:, :, 0]
material_ids_map = id_map[:, :, 2]
# Check that we have valid cell IDs (not all -2)
assert np.any(cell_id_map >= 0), "No valid cell IDs found in the id_map"
# Check that we have valid material IDs
assert np.any(material_ids_map >= 0), "No valid material IDs found in the id_map"
# Check that the expected cell IDs are present
expected_cell_ids = [10, 20] # Root cell, inner cell, outer cell
found_cell_ids = np.unique(cell_id_map[cell_id_map >= 0])
for cell_id in expected_cell_ids:
assert cell_id in found_cell_ids, f"Expected cell ID {cell_id} not found in id_map"
# Check that the expected material IDs are present
expected_material_ids = [1, 2, 3, 4, 5] # All materials defined above
found_material_ids = np.unique(material_ids_map[material_ids_map >= 0])
for mat_id in expected_material_ids:
assert mat_id in found_material_ids, f"Expected material ID {mat_id} not found in id_map"
# Test specific regions to verify lattice structure
# Check center of each lattice cell (should be inner cells)
# Lattice cell centers are at (-0.5, -0.5), (0.5, -0.5), (-0.5, 0.5), (0.5, 0.5)
# With 200x200 pixels over 2x2 units, each pixel is 0.01 units
# Bottom-left lattice cell center (should be inner cell 10)
bl_cell, bl_instance, bl_material = id_map[-50, 50]
assert bl_cell == 10, f"Expected cell ID 10 at bottom-left center, got {bl_cell}"
assert bl_instance == 0, f"Expected cell instance 0 at bottom-left center, got {bl_instance}"
assert bl_material == 1, f"Expected material ID 1 at bottom-left center, got {bl_material}"
# Bottom-right lattice cell center (should be inner cell 10)
br_cell, br_instance, br_material = id_map[-50, 150]
assert br_cell == 10, f"Expected cell ID 10 at bottom-right center, got {br_cell}"
assert br_instance == 1, f"Expected cell instance 1 at bottom-right center, got {br_instance}"
assert br_material == 2, f"Expected material ID 2 at bottom-right center, got {br_material}"
# Top-left lattice cell center (should be inner cell 10)
tl_cell, tl_instance, tl_material = id_map[-150, 50]
assert tl_cell == 10, f"Expected cell ID 10 at top-left center, got {tl_cell}"
assert tl_instance == 2, f"Expected cell instance 2 at top-left center, got {tl_instance}"
assert tl_material == 3, f"Expected material ID 3 at top-left center, got {tl_material}"
# Top-right lattice cell center (should be inner cell 10)
tr_cell, tr_instance, tr_material = id_map[-150, 150]
assert tr_cell == 10, f"Expected cell ID 10 at top-right center, got {tr_cell}"
assert tr_instance == 3, f"Expected cell instance 3 at top-right center, got {tr_instance}"
assert tr_material == 4, f"Expected material ID 4 at top-right center, got {tr_material}"
# Check that the model is properly finalized after id_map call
assert not model.is_initialized, "Model should be finalized after id_map call"
# Check that the values at the corners are correctly set as the outer cell and material
bl_cell, bl_instance, bl_material = id_map[-1, 0]
assert bl_cell == 20, f"Expected cell ID 20 at bottom-left corner, got {bl_cell}"
assert bl_instance == 0, f"Expected cell instance 0 at bottom-left corner, got {bl_instance}"
assert bl_material == 5, f"Expected material ID 5 at bottom-left corner, got {bl_material}"
br_cell, br_instance, br_material = id_map[-1, -1]
assert br_cell == 20, f"Expected cell ID 20 at bottom-right corner, got {br_cell}"
assert br_instance == 1, f"Expected cell instance 1 at bottom-right corner, got {br_instance}"
assert br_material == 5, f"Expected material ID 5 at bottom-right corner, got {br_material}"
tl_cell, tl_instance, tl_material = id_map[0, 0]
assert tl_cell == 20, f"Expected cell ID 20 at top-left corner, got {tl_cell}"
assert tl_instance == 2, f"Expected cell instance 2 at top-left corner, got {tl_instance}"
assert tl_material == 5, f"Expected material ID 5 at top-left corner, got {tl_material}"
tr_cell, tr_instance, tr_material = id_map[0, -1]
assert tr_cell == 20, f"Expected cell ID 20 at top-right corner, got {tr_cell}"
assert tr_instance == 3, f"Expected cell instance 3 at top-right corner, got {tr_instance}"
assert tr_material == 5, f"Expected material ID 5 at top-right corner, got {tr_material}"