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

using id_map in model.plot for more efficient plotting (#3678)

Browse source
This commit is contained in:
Jonathan Shimwell 2025-12-17 21:44:53 +01:00 • committed by GitHub
parent d118356638
commit e0eb91b955
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 211 additions and 108 deletions
Download patch file Download diff file Expand all files Collapse all files

195
openmc/model/model.py
View file

@ -23,7 +23,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, _BASIS_INDICES
from openmc.plots import add_plot_params, _BASIS_INDICES, id_map_to_rgb
from openmc.utility_funcs import change_directory
@ -1114,13 +1114,12 @@ class Model:
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,
show_overlaps: bool = False,
overlap_color: Sequence[int] | str | None = None,
overlap_color: Sequence[int] | str = (255, 0, 0),
n_samples: int | None = None,
plane_tolerance: float = 1.,
legend_kwargs: dict | None = None,
@ -1132,7 +1131,6 @@ class Model:
.. versionadded:: 0.15.1
"""
import matplotlib.image as mpimg
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
@ -1162,125 +1160,108 @@ class Model:
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]
# Determine whether any materials contains macroscopic data and if so,
# set energy mode accordingly
_energy_mode = self.settings._energy_mode
for mat in self.geometry.get_all_materials().values():
if mat._macroscopic is not None:
self.settings.energy_mode = 'multi-group'
break
# Get ID map from the C API
id_map = self.id_map(
origin=origin,
width=width,
pixels=pixels,
basis=basis,
color_overlaps=show_overlaps
)
with TemporaryDirectory() as tmpdir:
_plot_seed = self.settings.plot_seed
if seed is not None:
self.settings.plot_seed = seed
# Create plot object matching passed arguments
# Generate colors if not provided
if colors is None and seed is not None:
# Use the colorize method to generate random colors
plot = openmc.SlicePlot()
plot.origin = origin
plot.width = width
plot.pixels = pixels
plot.basis = basis
plot.color_by = color_by
plot.show_overlaps = show_overlaps
if overlap_color is not None:
plot.overlap_color = overlap_color
if colors is not None:
plot.colors = colors
self.plots.append(plot)
plot.colorize(self.geometry, seed=seed)
colors = plot.colors
# Run OpenMC in geometry plotting mode
self.plot_geometry(False, cwd=tmpdir, openmc_exec=openmc_exec)
# Convert ID map to RGB image
img = id_map_to_rgb(
id_map=id_map,
color_by=color_by,
colors=colors,
overlap_color=overlap_color
)
# Undo changes to model
self.plots.pop()
self.settings._plot_seed = _plot_seed
self.settings._energy_mode = _energy_mode
# 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_px = pixels[0]*px/(params.right - params.left)
height_px = pixels[1]*px/(params.top - params.bottom)
fig.set_size_inches(width_px, height_px)
# 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))
if outline:
# Combine R, G, B values into a single int for contour detection
rgb = (img * 256).astype(int)
image_value = (rgb[..., 0] << 16) + \
(rgb[..., 1] << 8) + (rgb[..., 2])
# 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)
# 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')
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",
levels=np.unique(image_value),
extent=(x_min, x_max, y_min, y_max),
**contour_kwargs
)
# If only showing outline, set the axis limits and aspect explicitly
if outline == 'only':
axes.set_xlim(x_min, x_max)
axes.set_ylim(y_min, y_max)
axes.set_aspect('equal')
# 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')
# Add legend showing which colors represent which material or cell
if legend:
if colors is None or len(colors) == 0:
raise ValueError("Must pass 'colors' dictionary if you "
"are adding a legend via legend=True.")
axes.contour(
image_value,
origin="upper",
levels=np.unique(image_value),
extent=(x_min, x_max, y_min, y_max),
**contour_kwargs
)
if color_by == "cell":
expected_key_type = openmc.Cell
else:
expected_key_type = openmc.Material
# 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.")
patches = []
for key, color in 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")
if color_by == "cell":
expected_key_type = openmc.Cell
# 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
if len(color) == 3 and not isinstance(color, str):
scaled_color = (
color[0]/255, color[1]/255, color[2]/255)
else:
expected_key_type = openmc.Material
scaled_color = color
patches = []
for key, color in plot.colors.items():
key_patch = mpatches.Patch(color=scaled_color, label=label)
patches.append(key_patch)
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)
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

82
openmc/plots.py
View file

@ -1,4 +1,4 @@
from collections.abc import Iterable, Mapping
from collections.abc import Iterable, Mapping, Sequence
from numbers import Integral, Real
from pathlib import Path
from textwrap import dedent
@ -355,6 +355,86 @@ def voxel_to_vtk(voxel_file: PathLike, output: PathLike = 'plot.vti'):
return output
def id_map_to_rgb(
id_map: np.ndarray,
color_by: str = 'cell',
colors: dict | None = None,
overlap_color: Sequence[int] | str = (255, 0, 0)
) -> np.ndarray:
"""Convert ID map array to RGB image array.
Parameters
----------
id_map : numpy.ndarray
Array with shape (v_pixels, h_pixels, 3) containing cell IDs,
cell instances, and material IDs
color_by : {'cell', 'material'}
Whether to color by cell or material
colors : dict, optional
Dictionary mapping cells/materials to colors
overlap_color : sequence of int or str, optional
Color to use for overlaps. Defaults to red (255, 0, 0).
Returns
-------
numpy.ndarray
RGB image array with shape (v_pixels, h_pixels, 3) with values
in range [0, 1] for matplotlib
"""
# Initialize RGB array with white background (values between 0 and 1 for matplotlib)
img = np.ones(id_map.shape, dtype=float)
# Get the appropriate index based on color_by
if color_by == 'cell':
id_index = 0 # Cell IDs are in the first channel
elif color_by == 'material':
id_index = 2 # Material IDs are in the third channel
else:
raise ValueError("color_by must be either 'cell' or 'material'")
# Get all unique IDs in the plot
unique_ids = np.unique(id_map[:, :, id_index])
# Generate default colors if not provided
if colors is None:
colors = {}
# Convert colors dict to use IDs as keys
color_map = {}
for key, color in colors.items():
if isinstance(key, (openmc.Cell, openmc.Material)):
color_map[key.id] = color
else:
color_map[key] = color
# Generate random colors for IDs not in color_map
rng = np.random.RandomState(1)
for uid in unique_ids:
if uid > 0 and uid not in color_map:
color_map[uid] = rng.randint(0, 256, (3,))
# Apply colors to each pixel
for uid in unique_ids:
if uid == -1: # Background/void
continue
elif uid == -3: # Overlap (only present if color_overlaps was True)
if isinstance(overlap_color, str):
rgb = _SVG_COLORS[overlap_color.lower()]
else:
rgb = overlap_color
mask = id_map[:, :, id_index] == uid
img[mask] = np.array(rgb) / 255.0
elif uid in color_map:
color = color_map[uid]
if isinstance(color, str):
rgb = _SVG_COLORS[color.lower()]
else:
rgb = color
mask = id_map[:, :, id_index] == uid
img[mask] = np.array(rgb) / 255.0
return img
class PlotBase(IDManagerMixin):
"""
Parameters

42
tests/unit_tests/test_model.py
View file

@ -7,6 +7,7 @@ import pytest
import openmc
import openmc.lib
from openmc.plots import id_map_to_rgb
@pytest.fixture(scope='function')
@ -996,3 +997,44 @@ def test_keff_search(run_in_tmpdir):
# Check that total_batches property works
assert result.total_batches == sum(result.batches)
assert result.total_batches > 0
def test_id_map_to_rgb():
"""Test conversion of ID map to RGB image array."""
# Create a simple model
mat = openmc.Material()
mat.set_density('g/cm3', 1.0)
mat.add_nuclide('Li7', 1.0)
sphere = openmc.Sphere(r=5.0, boundary_type='vacuum')
cell = openmc.Cell(fill=mat, region=-sphere)
geometry = openmc.Geometry([cell])
settings = openmc.Settings(
batches=10, particles=100, run_mode='fixed source'
)
model = openmc.Model(geometry, settings=settings)
id_data = np.zeros((10, 10, 3), dtype=np.int32)
id_data[:, :, 0] = cell.id # Cell IDs
id_data[:, :, 2] = mat.id # Material IDs
# Test color_by with default colors
for color_by in ['cell', 'material']:
rgb = id_map_to_rgb(id_data, color_by=color_by)
assert rgb.shape == (10, 10, 3)
assert rgb.dtype == float
assert np.all((rgb >= 0) & (rgb <= 1)) # RGB values in [0, 1]
# Test with custom colors
colors = {cell.id: (255, 0, 0)} # Red
rgb_custom = id_map_to_rgb(id_data, color_by='cell', colors=colors)
assert np.allclose(rgb_custom, [1.0, 0.0, 0.0]) # All pixels should be red
# Test with overlaps
id_data_overlap = id_data.copy()
id_data_overlap[5:, 5:, 0] = -3 # Mark some pixels as overlaps
rgb_overlap = id_map_to_rgb(
id_data_overlap, overlap_color=(0, 255, 0)
)
# Check that overlap region is green
assert np.allclose(rgb_overlap[5:, 5:], [0.0, 1.0, 0.0])