Adding regularmesh tally plot

[shimwell]

Description

Plotting mesh tallies along with the outline geometry in one function call.
This would be very helpful as a mesh tally plot is a very common output for posters, presentations papers.
The function allows lots of options for scaling, log plots, colorbars with args and lots of customizable options for imshow and contour via kwarg
testing of the plot includes some axis attribute checks and checking function runs but I've also been writting scripts to test the geometry like the one below

Checklist

• I have performed a self-review of my own code
  - [ ] I have run clang-format on any C++ source files (if applicable)
• 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)

Example script to show this function in use

from matplotlib.colors import LogNorm
import matplotlib.pyplot as plt
import openmc

mat1 = openmc.Material()
mat1.add_element("Al", 1.0)
mat1.set_density("g/cm3", 7)

surf_1 = openmc.XPlane(x0=-10)
surf_2 = openmc.XPlane(x0=15)
surf_3 = openmc.YPlane(y0=-5)
surf_4 = openmc.YPlane(y0=20)
surf_5 = openmc.ZPlane(z0=-25)
surf_6 = openmc.ZPlane(z0=25)

region1 = +surf_1 & -surf_2 & +surf_3 & -surf_4 & +surf_5 & -surf_6
cell1 = openmc.Cell(region=region1, fill=mat1)

surf_7 = openmc.XPlane(x0=-35, boundary_type="vacuum")
surf_8 = openmc.XPlane(x0=35, boundary_type="vacuum")
surf_9 = openmc.YPlane(y0=-35, boundary_type="vacuum")
surf_10 = openmc.YPlane(y0=35, boundary_type="vacuum")
surf_11 = openmc.ZPlane(z0=-35, boundary_type="vacuum")
surf_12 = openmc.ZPlane(z0=35, boundary_type="vacuum")

region2 = +surf_7 & -surf_8 & +surf_9 & -surf_10 & +surf_11 & -surf_12
cell2 = openmc.Cell(region=region2 & ~region1)

my_geometry = openmc.Geometry([cell1, cell2])
ww_r_mesh = openmc.RegularMesh().from_domain(my_geometry)
ww_r_mesh_filter = openmc.MeshFilter(ww_r_mesh)

n_flux_r_mesh_tally = openmc.Tally(name="flux_on_r_mesh")
n_flux_r_mesh_tally.filters = [ww_r_mesh_filter]
n_flux_r_mesh_tally.scores = ["flux", "heating"]

tallies = openmc.Tallies([n_flux_r_mesh_tally])

my_materials = openmc.Materials([mat1])

my_settings = openmc.Settings()
my_settings.batches = 2
my_settings.inactive = 0
my_settings.particles = 100000
my_settings.run_mode = "fixed source"

source = openmc.IndependentSource()
source.angle = openmc.stats.Isotropic()
source.energy = openmc.stats.Discrete([14e6], [1])
source.space = openmc.stats.Point(region1.bounding_box.center)
my_settings.source = source

model = openmc.Model(my_geometry, my_materials, my_settings, tallies)
filename = model.run()

with openmc.StatePoint(filename) as statepoint:
    neutron_tally = statepoint.get_tally(name="flux_on_r_mesh")

for score in ["flux", "heating"]:
    max_v = max(neutron_tally.get_slice(scores=[score]).mean.flatten())
    min_v = min(neutron_tally.get_slice(scores=[score]).mean.flatten()) + 0.001
    # + 0.001 avoids log plot with zero values

    plot = openmc.plot_mesh_tally(
        tally=neutron_tally,
        basis="xy",
        score=score,
        axis_units="cm",
        value="mean",
        outline=True,
        geometry=my_geometry,
        outline_by="cell",
        norm=LogNorm(vmin=min_v, vmax=max_v),
    )
    plot.figure.savefig(f"xy-{score}-neutron.png")

    plot = openmc.plot_mesh_tally(
        tally=neutron_tally,
        basis="xz",
        score=score,
        axis_units="cm",
        value="mean",
        outline=True,
        geometry=my_geometry,
        outline_by="cell",
        norm=LogNorm(vmin=min_v, vmax=max_v),
        cmap="gray",
        outline_kwargs={"color": "red"},
    )
    plot.figure.savefig(f"xz-{score}-neutron.png")

    plot = openmc.plot_mesh_tally(
        tally=neutron_tally,
        basis="yz",
        score=score,
        axis_units="cm",
        value="mean",
        outline=True,
        geometry=my_geometry,
        outline_by="cell",
        norm=LogNorm(vmin=min_v, vmax=max_v),
        cmap="Blues",
    )
    plot.title.set_text('custom title')
    plot.figure.savefig(f"yz-{score}-neutron.png")

I have another example script that makes the FLF logo over here if you want to see more test scripts

[paulromano]

I'm not sure how I feel about this one. When it comes to producing plots, my general thought is that for users who want to produce something easily, we should enable that through GUI applications like openmc-plotter and openmc-plot. Power users will likely want more customization over their plots, so the effort for them to figure out a function like this may not be worth it.

[shimwell]

Thanks for considering this PR @paulromano . Looks like it is not necessary to add this code to the openmc repository.

I still have a need for this type of functionality so I have created a small python package just in case it is useful for anyone else.

[paulromano]

Thanks for understanding @shimwell. I just gave you a ⭐ on your mini package 😄