Adding source option to plot

[shimwell]

Description

A while back there was a slack conversation where we talked about adding source plotting to the geometry plot.
@pshriwise recently put in a PR for adding source plotting to the openmc-plotter package.

So I made this PR that adds source plotting to the openmc python API for those users who might want the source term plotting on the geometry from the Python API.

Here are some examples of plots which have a source

import openmc
import matplotlib.pylab as plt

outer_sphere_surface = openmc.Sphere(r=600, boundary_type='vacuum')
blanket_region = -outer_sphere_surface  
blanket_cell = openmc.Cell(region=blanket_region)
geometry = openmc.Geometry([blanket_cell])  
settings = openmc.Settings()
settings.particles = 100
settings.batches = 100
model = openmc.Model(geometry, None,settings)
plot = model.plot(n_samples=1, plane_tolerance= 1.)
plot.figure.savefig('point_xz.png', bbox_inches='tight')

outer_sphere_surface = openmc.Sphere(r=600, boundary_type='vacuum')
blanket_region = -outer_sphere_surface  
blanket_cell = openmc.Cell(region=blanket_region)
geometry = openmc.Geometry([blanket_cell])
my_source = openmc.IndependentSource()
z_values = openmc.stats.Discrete([4.9], [1])
angle = openmc.stats.Uniform(a=0., b=2* 3.14159265359)
radius = openmc.stats.Discrete([250], [1])
my_source.space = openmc.stats.CylindricalIndependent(r=radius, phi=angle, z=z_values)
settings = openmc.Settings()
settings.source = my_source
settings.particles = 100
settings.batches = 100
model = openmc.Model(geometry, None,settings)
plot = model.plot(n_samples=1000, plane_tolerance=5., basis='xz')
plot.figure.savefig('ring_xz.png', bbox_inches='tight')

plot = model.plot(n_samples=50, plane_tolerance=5., basis='xy')
plot.figure.savefig('ring_xy.png', bbox_inches='tight')

from openmc_plasma_source import TokamakSource
my_sources = TokamakSource(
    elongation=1.557,
    ion_density_centre=1.09e20,
    ion_density_peaking_factor=1,
    ion_density_pedestal=1.09e20,
    ion_density_separatrix=3e19,
    ion_temperature_centre=45.9,
    ion_temperature_peaking_factor=8.06,
    ion_temperature_pedestal=6.09,
    ion_temperature_separatrix=0.1,
    major_radius=9.06,
    minor_radius=2.92258,
    pedestal_radius=0.8 * 2.92258,
    mode="H",
    shafranov_factor=0.44789,
    triangularity=0.270,
    ion_temperature_beta=6,
    sample_size=50000,  # the number of individual sources to use to make a combined source
    angles=(0, 0.01)  # angle in radians
).make_openmc_sources()
settings = openmc.Settings()
settings.particles = 1
settings.batches = 1
settings.source = my_sources
materials = openmc.Materials()
sph = openmc.Sphere(r=5, x0=9.06,boundary_type="vacuum")
cell = openmc.Cell(region=-sph)
geometry = openmc.Geometry([cell])
model = openmc.Model(geometry, materials, settings)
plot = model.plot(n_samples=1500, plane_tolerance=5., basis='xz', pixels=400000, source_kwargs={'marker':'o','color':'red'})
plot.figure.savefig('toko.png', bbox_inches='tight')

tagging @Asureda for your interest

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]

Added some tests that get the coordinates of the scatter plot and compare with source location for different basis and with different plane_tolerances

[shimwell]

Sorry to nudge, is anyone able to review this one. I'm keen to see if it gets accepted or not.

[jon-proximafusion]

Just keeping this branch upto date with develop. I'm making use of this branch for plotting stellerator sources and finding it quite useful.

[paulromano]

Sorry it took so long to get this one resolved @shimwell @jon-proximafusion! I just made a few updates to your branch. Most importantly, the sample_external_source method in Model executes in a temporary directory so as not to mess with any local files. If you're happy with the changes, let me know and we can get this merged in!

[shimwell]

oh wow, you have made my day @paulromano
I had forgotten about this PR but am totally delighted to see it merged in.

[shimwell]

Thanks for the improvements, sorry it needed so many changes. I see the temporary dir and yes that is certainly a good idea to avoid overwriting local xml files. The changes look great. I shall enable auto merge so we don't have to come back.