added slice_of_data method for easy RegularMesh plotting

[shimwell]

This PR proposes a small helper method for users wanting to plot a slice of data on a regular mesh.

I must have written this code three times already for different mesh plotting apps but recently another community member had need for something like this so I have finally got around to making a PR.

The view directions follow the paraview convention of axis orientation

This slice_of_data method helps people wanting to plot a matplotlib imshow of the mesh tally.

If this method is accepted then I also have the same functionality for the CylindricalMesh class that I would be keen to add.

If this method is welcome then I would also propose to build another method called something like plot_slice that makes use of this slice_of_data method and the mesh bounding box to define the matplotlib imshow extent. But I wanted to take it one step at a time.

I have a script that produces some test plots of a geometry that was made to check the axis are potted correctly.

All the 6 axis directions are produced by the test script but here is an example output of test script

test script with geometry designed to check plots are correct

# creates a simple CSG geometry with different size spheres on the positive x,
# y, and z axis. Performs a mesh tally on the geometry and plots slices of the
# resulting mesh tally with matplotlib and making use of the proposed
# slice_of_data function

import os

import matplotlib.pyplot as plt
import openmc

breeder_material = openmc.Material()
breeder_material.add_element("Li", 1.0)
breeder_material.set_density("g/cm3", 0.01)

my_materials = openmc.Materials([breeder_material])

# surfaces
central_sphere_surf = openmc.Sphere(r=0.001, x0=0.0, y0=0.0, z0=0.0)
x_positve_sphere_surf = openmc.Sphere(r=1, x0=5., y0=0.0, z0=0.0)
y_positve_sphere_surf = openmc.Sphere(r=2, x0=0.0, y0=5, z0=0.0)
z_positve_sphere_surf = openmc.Sphere(r=3, x0=0.0, y0=0, z0=5)
encasing_sphere_surf = openmc.Sphere(r=10, x0=0.0, y0=0.0, z0=0.0, boundary_type='vacuum')

# regions
central_sphere_region = -central_sphere_surf
x_positve_sphere_region = -x_positve_sphere_surf
y_positve_sphere_region = -y_positve_sphere_surf
z_positve_sphere_region = -z_positve_sphere_surf
encasing_sphere_region = -encasing_sphere_surf & \
                ~central_sphere_region & \
                ~x_positve_sphere_region & \
                ~y_positve_sphere_region & \
                ~z_positve_sphere_region

# cells
central_sphere_cell = openmc.Cell(region=central_sphere_region)
central_sphere_cell.fill = breeder_material

x_positve_sphere_cell = openmc.Cell(region=x_positve_sphere_region)
x_positve_sphere_cell.fill = breeder_material
y_positve_sphere_cell = openmc.Cell(region=y_positve_sphere_region)
y_positve_sphere_cell.fill = breeder_material
z_positve_sphere_cell = openmc.Cell(region=z_positve_sphere_region)
z_positve_sphere_cell.fill = breeder_material

encasing_sphere_cell = openmc.Cell(region=encasing_sphere_region)

universe = openmc.Universe(
    cells=[
        central_sphere_cell,
        x_positve_sphere_cell,
        y_positve_sphere_cell,
        z_positve_sphere_cell,
        encasing_sphere_cell,
    ]
)
my_geometry = openmc.Geometry(universe)

my_settings = openmc.Settings()
my_settings.batches = 3
my_settings.particles = 50000
my_settings.run_mode = 'fixed source'

my_source = openmc.Source()
my_source.angle = openmc.stats.Isotropic()
# enough energy to perform n,2n in Be but not n,xt in Be
my_source.energy = openmc.stats.Discrete([3e6], [1])
my_source.space = openmc.stats.Point((0, 0, 0))


my_settings.source = my_source
dimension=(50,70,100)
mesh = openmc.RegularMesh().from_domain(
    domain=my_geometry,
    dimension=dimension
)

my_tallies = openmc.Tallies()
# Create mesh filter for tally
mesh_filter = openmc.MeshFilter(mesh)

mesh_tally = openmc.Tally(name='n_xt_on_mesh')
mesh_tally.filters = [mesh_filter]
mesh_tally.scores = ['(n,Xt)']  # where X is a wildcard
my_tallies.append(mesh_tally)

model = openmc.model.Model(my_geometry, my_materials, my_settings, my_tallies)
sp_filename = model.run()

# loads up the output file from the simulation
statepoint = openmc.StatePoint(sp_filename)

# extracts the mesh tally by name
tally = statepoint.get_tally(name='n_xt_on_mesh')

# plots a imshow for each view direction
# view directions are same as paraview
for view_direction in ["x", "-x", "y", "-y", "z", "-z"]:
    if 'x' in view_direction:
        slice_index = int(dimension[0]/2)
    if 'y' in view_direction:
        slice_index = int(dimension[1]/2)
    if 'z' in view_direction:
        slice_index = int(dimension[2]/2)
    for tally in [my_tbr_tally]:
        im = mesh.slice_of_data(
            dataset=tally.mean,
            view_direction=view_direction,
            slice_index=slice_index # mid mesh
        )
        plt.cla()
        plt.clf()
        fig, ax = plt.subplots()
        plot=plt.imshow(im, interpolation='None')
        fig.colorbar(plot, orientation='vertical')
        plt.title(f'view_direction {view_direction}')
        plt.savefig(f'view_direction_{view_direction}.png')

# plots a sweep through the slice indexes of the x axis
for slice_index in range(0,50):
    for tally in [my_tbr_tally]:
        im = mesh.slice_of_data(
            dataset=tally.mean,
            view_direction='x',
            slice_index=slice_index
        )
        plt.cla()
        plt.clf()
        plt.imshow(im, interpolation='None')
        plt.title(f'view_direction x_slice_index_{slice_index}')
        plt.savefig(f'view_direction_x_slice_index_{str(slice_index).zfill(3)}.png')

# makes gif animation
os.system('convert -delay 10 -loop 0 view_direction_x_slice_index_*.png x_sweep.gif')

[shimwell]

ok I couldn't resist adding the plotting method that makes use of this data slice. I would be keen to keep them separate as some users might want to pass the data into a different plotting backend (e.g. plotly).

Updated script is a bit more concise as the plotting is now done within the openmc source code. I've also added the ability to pass in imshow kwargs much like the universe.plot and made use of the log scale arg in the updated example script.

# creates a simple CSG geometry with different size spheres on the positive x,
# y, and z axis. Performs a mesh tally on the geometry and plots slices of the
# resulting mesh tally with matplotlib and making use of the proposed
# plot_slice function

import os

import matplotlib.pyplot as plt
import openmc

breeder_material = openmc.Material()
breeder_material.add_element("Li", 1.0)
breeder_material.set_density("g/cm3", 0.01)

my_materials = openmc.Materials([breeder_material])

# surfaces
central_sphere_surf = openmc.Sphere(r=0.001, x0=0.0, y0=0.0, z0=0.0)
x_positve_sphere_surf = openmc.Sphere(r=1, x0=5., y0=0.0, z0=0.0)
y_positve_sphere_surf = openmc.Sphere(r=2, x0=0.0, y0=5, z0=0.0)
z_positve_sphere_surf = openmc.Sphere(r=3, x0=0.0, y0=0, z0=5)
encasing_sphere_surf = openmc.Sphere(r=10, x0=0.0, y0=0.0, z0=0.0, boundary_type='vacuum')

# regions
central_sphere_region = -central_sphere_surf
x_positve_sphere_region = -x_positve_sphere_surf
y_positve_sphere_region = -y_positve_sphere_surf
z_positve_sphere_region = -z_positve_sphere_surf
encasing_sphere_region = -encasing_sphere_surf & \
                ~central_sphere_region & \
                ~x_positve_sphere_region & \
                ~y_positve_sphere_region & \
                ~z_positve_sphere_region

# cells
central_sphere_cell = openmc.Cell(region=central_sphere_region)
central_sphere_cell.fill = breeder_material

x_positve_sphere_cell = openmc.Cell(region=x_positve_sphere_region)
x_positve_sphere_cell.fill = breeder_material
y_positve_sphere_cell = openmc.Cell(region=y_positve_sphere_region)
y_positve_sphere_cell.fill = breeder_material
z_positve_sphere_cell = openmc.Cell(region=z_positve_sphere_region)
z_positve_sphere_cell.fill = breeder_material

encasing_sphere_cell = openmc.Cell(region=encasing_sphere_region)

universe = openmc.Universe(
    cells=[
        central_sphere_cell,
        x_positve_sphere_cell,
        y_positve_sphere_cell,
        z_positve_sphere_cell,
        encasing_sphere_cell,
    ]
)
my_geometry = openmc.Geometry(universe)

my_settings = openmc.Settings()
my_settings.batches = 3
my_settings.particles = 50000
my_settings.run_mode = 'fixed source'

my_source = openmc.Source()
my_source.angle = openmc.stats.Isotropic()
# enough energy to perform n,2n in Be but not n,xt in Be
my_source.energy = openmc.stats.Discrete([3e6], [1])
my_source.space = openmc.stats.Point((0, 0, 0))


my_settings.source = my_source
dimension=(50,70,100)
mesh = openmc.RegularMesh().from_domain(
    domain=my_geometry,
    dimension=dimension
)

my_tallies = openmc.Tallies()
# Create mesh filter for tally
mesh_filter = openmc.MeshFilter(mesh)

mesh_tally = openmc.Tally(name='n_xt_on_mesh')
mesh_tally.filters = [mesh_filter]
mesh_tally.scores = ['(n,Xt)']  # where X is a wildcard
my_tallies.append(mesh_tally)

model = openmc.model.Model(my_geometry, my_materials, my_settings, my_tallies)
sp_filename = model.run()

# loads up the output file from the simulation
statepoint = openmc.StatePoint(sp_filename)

# extracts the mesh tally by name
tally = statepoint.get_tally(name='n_xt_on_mesh')

# plots a imshow for each view direction
# view directions are same as paraview
from matplotlib.colors import LogNorm
for view_direction in ["x", "-x", "y", "-y", "z", "-z"]:
    mesh.plot_slice(
        dataset=tally.mean,
        view_direction=view_direction,
        norm=LogNorm(),
    )
    plt.savefig(f'view_direction_{view_direction}.png')

# makes gif animation
os.system('convert -delay 30 -loop 0 view_direction_*.png views.gif')

[shimwell]

Just to mention this mesh plotting method can also be used to plot a slice of weight windows bounds on the mesh they were recorded as well as tally.mean or tally.std_dev data

[emaartensson]

Hi @shimwell

Great work. I've personally been doing exactly what you have in terms of plotting slices, mostly for quick checks of results. The annoying thing is that reshaping the data to the mesh returns it with the wrong axes. The addition of both slice_of_data and plot_slice would make this a lot easier so I definitely support it!

[shimwell]

added a normalization argument like we have in the mesh.write_data_to_vtk method

[shimwell]

I just noticed the plotter also has similar functions for rotating and slicing the regular mesh tally

https://github.com/openmc-dev/plotter/blob/4f353ccee8645caf04fb6c5d71cd6136183d50d5/openmc_plotter/plotmodel.py#L636

If this was available in OpenMC then the plotter could make use of it instead of implementing the rotation and slicing itself. This movement of code would have the added benefit of allowing API access for the non GUI users as well.

I think the plotter code is perhaps nicer than this PR, so perhaps a bit of refactoring can be done on this PR

[shimwell]

I've now put this slice mesh data functionality in regular_mesh_plotter, I've also got some similar slice materials ID and slice cell ID in openmc-geometry-plot. With these combined packages I can make geometry or mesh tally plots with outlines. I'm happy to move that functionality over to openmc if it is ever needed. In general I think this sort of functionality should live in OpenMC so that is can be used but the plotter and users directly. But we have quite a lot of PRs open and this one was a bit of a mess anyway so I shall close it for now but happy to have another go if we want this functionality in openmc