diff --git a/openmc/model/model.py b/openmc/model/model.py index a9aaa481d..6e4c1c585 100644 --- a/openmc/model/model.py +++ b/openmc/model/model.py @@ -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 diff --git a/openmc/plots.py b/openmc/plots.py index cb722abc6..8b67d5cac 100644 --- a/openmc/plots.py +++ b/openmc/plots.py @@ -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 diff --git a/tests/unit_tests/test_model.py b/tests/unit_tests/test_model.py index d553af53c..3846ba4fb 100644 --- a/tests/unit_tests/test_model.py +++ b/tests/unit_tests/test_model.py @@ -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])