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

• [x] I have performed a self-review of my own code
• [x] I have followed the style guidelines for Python source files (if applicable)
• [x] I have made corresponding changes to the documentation (if applicable)
• [x] 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.