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

Author: shimwell

openmc/model/model.py

@@ -23,7 +23,7 @@
 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 @@ def plot(
         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 @@ def plot(
 
         .. versionadded:: 0.15.1
         """
-        import matplotlib.image as mpimg
         import matplotlib.patches as mpatches
         import matplotlib.pyplot as plt
 
@@ -1162,125 +1160,108 @@ def plot(
         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
-
-        with TemporaryDirectory() as tmpdir:
-            _plot_seed = self.settings.plot_seed
-            if seed is not None:
-                self.settings.plot_seed = seed
+        # Get ID map from the C API
+        id_map = self.id_map(
+            origin=origin,
+            width=width,
+            pixels=pixels,
+            basis=basis,
+            color_overlaps=show_overlaps
+        )
 
-            # 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)
-
-            # Run OpenMC in geometry plotting mode
-            self.plot_geometry(False, cwd=tmpdir, openmc_exec=openmc_exec)
-
-            # Undo changes to model
-            self.plots.pop()
-            self.settings._plot_seed = _plot_seed
-            self.settings._energy_mode = _energy_mode
-
-            # 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))
-
-            # 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)
-
-            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])
-
-                # 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')
-
-                axes.contour(
-                    image_value,
-                    origin="upper",
-                    levels=np.unique(image_value),
-                    extent=(x_min, x_max, y_min, y_max),
-                    **contour_kwargs
-                )
-
-            # 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.")
+            plot.colorize(self.geometry, seed=seed)
+            colors = plot.colors
+
+        # 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
+        )
 
-                if color_by == "cell":
-                    expected_key_type = openmc.Cell
+        # 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)
+
+        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])
+
+            # 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')
+
+            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')
+
+        # 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.")
+
+            if color_by == "cell":
+                expected_key_type = openmc.Cell
+            else:
+                expected_key_type = openmc.Material
+
+            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")
+
+                # 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
-
-                patches = []
-                for key, color in plot.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")
-
-                    # 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)
+                    scaled_color = color
 
-                axes.legend(handles=patches, **legend_kwargs)
+                key_patch = mpatches.Patch(color=scaled_color, label=label)
+                patches.append(key_patch)
 
-            # 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

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

tests/unit_tests/test_model.py

@@ -7,6 +7,7 @@
 
 import openmc
 import openmc.lib
+from openmc.plots import id_map_to_rgb
 
 
 @pytest.fixture(scope='function')
@@ -996,3 +997,44 @@ def modify_radius(radius):
     # 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])