using id_map in model.plot for more efficient plotting

[shimwell]

Description

This PR makes use of the model.id_map when plotting. This avoids writing the png/ppm file and then reading the image back in with mpimg.imread. I believe this is therefore more efficient than the current model.plot. Initially it looks like a lot of changes but many changes are just removing indentation which is no longer needed as the temp dir is not needed as we don't write files anymore.

I also split out some logic for converting id_maps to RGB maps which I could then test separately.

Previous PRs that made this possible
Following on from the excellent id_map PR #3481 that @pshriwise we recently added the ability to include overlaps in the id_map with #3669

Fixes # (issue)

Checklist

• I have performed a self-review of my own code
• I have followed the style guidelines for Python source files (if applicable)
• I have made corresponding changes to the documentation (if applicable)
• I have added tests that prove my fix is effective or that my feature works (if applicable)

[shimwell]

I have also been testing all the args with a local script I made, I attach it here in case it helps with the review process

#!/usr/bin/env python3
"""Test the refactored plot() method using id_map()"""

import openmc
import numpy as np

mat_air = openmc.Material(name="Air")
mat_air.add_element("N", 0.784431)
mat_air.add_element("O", 0.210748)
mat_air.add_element("Ar", 0.0046)
mat_air.set_density("g/cc", 0.001205)

mat_concrete = openmc.Material()
mat_concrete.add_element("H",0.168759)
mat_concrete.add_element("C",0.001416)
mat_concrete.add_element("O",0.562524)
mat_concrete.add_element("Na",0.011838)
mat_concrete.add_element("Mg",0.0014)
mat_concrete.add_element("Al",0.021354)
mat_concrete.add_element("Si",0.204115)
mat_concrete.add_element("K",0.005656)
mat_concrete.add_element("Ca",0.018674)
mat_concrete.add_element("Fe",0.00426)
mat_concrete.set_density("g/cm3", 2.3)

materials = openmc.Materials([mat_air, mat_concrete])

width_a = 100
width_b = 100
width_c = 500
width_d = 100
width_e = 100
width_f = 100
width_g = 100

depth_a = 100
depth_b = 100
depth_c = 700
depth_d = 600
depth_e = 100
depth_f = 100

height_j = 100
height_k = 500
height_l = 100

xplane_0 = openmc.XPlane(x0=0, boundary_type="vacuum")
xplane_1 = openmc.XPlane(x0=xplane_0.x0 + width_a)
xplane_2 = openmc.XPlane(x0=xplane_1.x0 + width_b)
xplane_3 = openmc.XPlane(x0=xplane_2.x0 + width_c)
xplane_4 = openmc.XPlane(x0=xplane_3.x0 + width_d)
xplane_5 = openmc.XPlane(x0=xplane_4.x0 + width_e)
xplane_6 = openmc.XPlane(x0=xplane_5.x0 + width_f)
xplane_7 = openmc.XPlane(x0=xplane_6.x0 + width_g, boundary_type="vacuum")

yplane_0 = openmc.YPlane(y0=0, boundary_type="vacuum")
yplane_1 = openmc.YPlane(y0=yplane_0.y0 + depth_a)
yplane_2 = openmc.YPlane(y0=yplane_1.y0 + depth_b)
yplane_3 = openmc.YPlane(y0=yplane_2.y0 + depth_c)
yplane_4 = openmc.YPlane(y0=yplane_3.y0 + depth_d)
yplane_5 = openmc.YPlane(y0=yplane_4.y0 + depth_e)
yplane_6 = openmc.YPlane(y0=yplane_5.y0 + depth_f, boundary_type="vacuum")

zplane_1 = openmc.ZPlane(z0=0, boundary_type="vacuum")
zplane_2 = openmc.ZPlane(z0=zplane_1.z0 + height_j)
zplane_3 = openmc.ZPlane(z0=zplane_2.z0 + height_k)
zplane_4 = openmc.ZPlane(z0=zplane_3.z0 + height_l, boundary_type="vacuum")

sphere = openmc.Sphere(r=50, x0=300, y0=400, z0=350)  # added to check overlap plotting

outside_left_region = +xplane_0 & -xplane_1 & +yplane_1 & -yplane_5 & +zplane_1 & -zplane_4
wall_left_region = +xplane_1 & -xplane_2 & +yplane_2 & -yplane_4 & +zplane_2 & -zplane_3
wall_right_region = +xplane_5 & -xplane_6 & +yplane_2 & -yplane_5 & +zplane_2 & -zplane_3
wall_top_region = +xplane_1 & -xplane_4 & +yplane_4 & -yplane_5 & +zplane_2 & -zplane_3
outside_top_region = +xplane_0 & -xplane_7 & +yplane_5 & -yplane_6 & +zplane_1 & -zplane_4
wall_bottom_region = +xplane_1 & -xplane_6 & +yplane_1 & -yplane_2 & +zplane_2 & -zplane_3
outside_bottom_region = +xplane_0 & -xplane_7 & +yplane_0 & -yplane_1 & +zplane_1 & -zplane_4
wall_middle_region = +xplane_3 & -xplane_4 & +yplane_3 & -yplane_4 & +zplane_2 & -zplane_3
outside_right_region = +xplane_6 & -xplane_7 & +yplane_1 & -yplane_5 & +zplane_1 & -zplane_4

room_region = +xplane_2 & -xplane_3 & +yplane_2 & -yplane_4 & +zplane_2 & -zplane_3
gap_region = +xplane_3 & -xplane_4 & +yplane_2 & -yplane_3 & +zplane_2 & -zplane_3
corridor_region = +xplane_4 & -xplane_5 & +yplane_2 & -yplane_5 & +zplane_2 & -zplane_3

roof_region = +xplane_1 & -xplane_6 & +yplane_1 & -yplane_5 & +zplane_1 & -zplane_2
floor_region = +xplane_1 & -xplane_6 & +yplane_1 & -yplane_5 & +zplane_3 & -zplane_4

outside_left_cell = openmc.Cell(region=outside_left_region, fill=mat_air)
outside_right_cell = openmc.Cell(region=outside_right_region, fill=mat_air)
outside_top_cell = openmc.Cell(region=outside_top_region, fill=mat_air)
outside_bottom_cell = openmc.Cell(region=outside_bottom_region, fill=mat_air)
wall_left_cell = openmc.Cell(region=wall_left_region, fill=mat_concrete)
wall_right_cell = openmc.Cell(region=wall_right_region, fill=mat_concrete)
wall_top_cell = openmc.Cell(region=wall_top_region, fill=mat_concrete)
wall_bottom_cell = openmc.Cell(region=wall_bottom_region, fill=mat_concrete)
wall_middle_cell = openmc.Cell(region=wall_middle_region, fill=mat_concrete)
room_cell = openmc.Cell(region=room_region, fill=mat_air)
gap_cell = openmc.Cell(region=gap_region, fill=mat_air)
corridor_cell = openmc.Cell(region=corridor_region, fill=mat_air)
overlapping_cell = openmc.Cell(region=-sphere, fill=mat_air)

roof_cell = openmc.Cell(region=roof_region, fill=mat_concrete)
floor_cell = openmc.Cell(region=floor_region, fill=mat_concrete)

geometry = openmc.Geometry(
    [
        overlapping_cell,
        outside_bottom_cell,
        outside_top_cell,
        outside_left_cell,
        outside_right_cell,
        wall_left_cell,
        wall_right_cell,
        wall_top_cell,
        wall_bottom_cell,
        wall_middle_cell,
        room_cell,
        gap_cell,
        corridor_cell,
        roof_cell,
        floor_cell,
    ]
)

settings = openmc.Settings(particles=100, batches=5)

# Get geometry bounding box center z value
bbox_center_z = geometry.bounding_box.center[2]

# Create ring source with radius 1m at center z height
ring_source = openmc.IndependentSource()
ring_source.space = openmc.stats.CylindricalIndependent(
    r=openmc.stats.Discrete([100.0], [1.0]),  # radius = 1m = 100cm
    phi=openmc.stats.Uniform(0, 2*np.pi),
    z=openmc.stats.Discrete([bbox_center_z], [1.0]),  # single z value at center
    origin=(600.0, 800.0, 0.0)
)
settings.source = ring_source

model = openmc.Model(geometry, settings=settings)

import matplotlib.pyplot as plt

# Test 1: Basic plot without passing axes
axes1 = model.plot(
    pixels=4000000,
    basis='xy',
    color_by='material',
    legend=True,
    outline=True,
    axis_units='m',
    colors={mat_air: 'lightblue', mat_concrete: 'gray'},
    show_overlaps=True,
    contour_kwargs={'linewidths': 5},
    n_samples=10
)

# Test 2: Pass in existing axes to verify the parameter works
fig, ax = plt.subplots(figsize=(10, 8))
axes2 = model.plot(
    pixels=4000000,
    basis='xy',
    color_by='material',
    legend=True,
    outline=True,
    axis_units='m',
    colors={mat_air: 'lightblue', mat_concrete: 'gray'},
    show_overlaps=True,
    axes=ax
)

plt.show()

[GuySten]

Looks good, @shimwell!

[pshriwise]

Nice refactor to reuse Model.id_map here @shimwell. Thank you!

Just a couple of quick questions from me but otherwise looks good.

[pshriwise]

Thanks for the responses, @shimwell! Hammered on this a bit and it worked well.

[pshriwise]

@GuySten I'll leave the merge to you since you had it lined up earlier today.