Merge pull request #39 from fusion-energy/change_from_extending_class…

…_to_functions

change package to user functions instead of extending openmc classes

.github/workflows/ci_with_install.yml

@@ -27,7 +27,7 @@ jobs:
       image: openmc/openmc:develop
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
 
       - name: install package
         run: |

README.md

@@ -68,7 +68,7 @@ A list of ```openmc.Particle``` objects can be obtained using ```model.sample_in
 
 ```python
 import openmc
-import openmc_source_plotter  # extents openmc.Model with sample_initial_particles method
+from openmc_source_plotter import sample_initial_particles
 
 settings = openmc.Settings()
 settings.particles = 1
@@ -83,7 +83,7 @@ geometry = openmc.Geometry([cell])
 
 model = openmc.Model(geometry, materials, settings)
 
-particles = model.sample_initial_particles(n_samples=10)
+particles = sample_initial_particles(this=model, n_samples=10)
 
 print(particles)
 >>>[<SourceParticle: neutron at E=1.440285e+07 eV>, <SourceParticle: neutron at E=1.397691e+07 eV>, <SourceParticle: neutron at E=1.393681e+07 eV>, <SourceParticle: neutron at E=1.470896e+07 eV>, <SourceParticle: neutron at E=1.460563e+07 eV>, <SourceParticle: neutron at E=1.420684e+07 eV>, <SourceParticle: neutron at E=1.413932e+07 eV>, <SourceParticle: neutron at E=1.412428e+07 eV>, <SourceParticle: neutron at E=1.464779e+07 eV>, <SourceParticle: neutron at E=1.391648e+07 eV>]

examples/example_gamma_spec_plot.py

@@ -1,5 +1,5 @@
 import openmc
-import openmc_source_plotter  # adds plot_gamma_emission plot to materials
+from openmc_source_plotter import plot_gamma_emission
 
 # this path will need changing to point to your chain file
 # openmc.config["chain_file"] = "chain-endf.xml"
@@ -12,6 +12,6 @@
 my_material.volume = 1  # must be set so number of atoms can be found
 
 # adds labels to the most active 3 gamma energies
-plt = my_material.plot_gamma_emission(label_top=3)
+plt = plot_gamma_emission(material=my_material, label_top=3)
 plt.xscale("log")  # modify axis from default settings
 plt.savefig("gamma_spec.png")

examples/example_get_particle_data.py

@@ -1,5 +1,5 @@
 import openmc
-import openmc_source_plotter  # overwrites the openmc.source method
+from openmc_source_plotter import sample_initial_particles
 
 # initialises a new source object
 my_source = openmc.Source()
@@ -14,6 +14,9 @@
 my_source.energy = openmc.stats.Discrete([14e6], [1])
 
 # gets the particle corrdiantes, energy and direction
-data = my_source.sample_initial_particles()
+particles = sample_initial_particles(my_source)
 
-print(data)
+print(particles)
+
+for particle in particles:
+    print(particle.E)

examples/example_plot_plasma_source_position.py

@@ -1,10 +1,8 @@
-from openmc_plasma_source import TokamakSource
+from openmc_plasma_source import tokamak_source
+from openmc_source_plotter import plot_source_position
+import openmc
 
-# openmc_plasma_source makes use of this package and
-# TokamakSource is a SourceWithPlotting object so it has
-# access to the plotting methods
-
-my_sources = TokamakSource(
+my_sources = tokamak_source(
     elongation=1.557,
     ion_density_centre=1.09e20,
     ion_density_peaking_factor=1,
@@ -23,12 +21,13 @@
     ion_temperature_beta=6,
     angles=(0, 3.14),  # makes a sector of 0 radians to 3.14 radians
     sample_size=100,  # reduces the number of samples from a default of 1000 to reduce plot time
-).make_openmc_sources()
+)
+
+settings = openmc.Settings()
+settings.Source = my_sources
 
 
 # plots the particle energy distribution
-plot = None
-for source in my_sources:
-    plot = source.plot_source_position(figure=plot)
+plot = plot_source_position(settings)
 
 plot.show()

examples/example_plot_source_direction.py

@@ -1,13 +1,13 @@
 import openmc
-import openmc_source_plotter  # overwrites the openmc.source method
+from openmc_source_plotter import plot_source_direction
 
 # initializes a new source object
-my_source = openmc.Source()
+my_source = openmc.IndependentSource()
 
 # sets the direction to isotropic
 my_source.angle = openmc.stats.Isotropic()
 
 # plots the particle energy distribution
-plot = my_source.plot_source_direction(n_samples=200)
+plot = plot_source_direction(this=my_source, n_samples=200)
 
 plot.show()

examples/example_plot_source_energy.py

@@ -1,13 +1,13 @@
 import openmc
-import openmc_source_plotter  # overwrites the openmc.source method
+from openmc_source_plotter import plot_source_energy
 
 # initialise a new source object
-my_source = openmc.Source()
+my_source = openmc.IndependentSource()
 
-# sets the energy distribution to a Muir distribution neutrons
-my_source.energy = openmc.stats.Muir(e0=14080000.0, m_rat=5.0, kt=20000.0)
+# sets the energy distribution to a muir distribution neutrons
+my_source.energy = openmc.stats.muir(e0=14080000.0, m_rat=5.0, kt=20000.0)
 
 # plots the particle energy distribution
-plot = my_source.plot_source_energy(n_samples=10000)
+plot = plot_source_energy(this=my_source, n_samples=10000)
 
 plot.show()

examples/example_plot_source_position.py

@@ -1,5 +1,5 @@
 import openmc
-import openmc_source_plotter  # overwrites the openmc.source method
+from openmc_source_plotter import plot_source_position
 
 # initialises a new source object
 my_source = openmc.Source()
@@ -20,6 +20,6 @@
 )
 
 # plots the particle energy distribution
-plot = my_source.plot_source_position()
+plot = plot_source_position(my_source)
 
 plot.show()

examples/example_plot_two_source_energies.py

@@ -1,19 +1,19 @@
 import openmc
-import openmc_source_plotter  # overwrites the openmc.source method
+from openmc_source_plotter import plot_source_energy
 
 # initialises a new source object
 my_dt_source = openmc.Source()
 
-# sets the energy distribution to a Muir distribution DT neutrons
-my_dt_source.energy = openmc.stats.Muir(e0=14080000.0, m_rat=5.0, kt=20000.0)
+# sets the energy distribution to a muir distribution DT neutrons
+my_dt_source.energy = openmc.stats.muir(e0=14080000.0, m_rat=5.0, kt=20000.0)
 
 # initialises a new source object
 my_dd_source = openmc.Source()
-# sets the energy distribution to a Muir distribution DD neutrons
-my_dd_source.energy = openmc.stats.Muir(e0=2080000.0, m_rat=2.0, kt=20000.0)
+# sets the energy distribution to a muir distribution DD neutrons
+my_dd_source.energy = openmc.stats.muir(e0=2080000.0, m_rat=2.0, kt=20000.0)
 
 # plots the particle energy distribution
-figure1 = my_dd_source.plot_source_energy(n_samples=10000)
-figure2 = my_dt_source.plot_source_energy(figure=figure1, n_samples=10000)
+figure1 = plot_source_energy(this=my_dd_source, n_samples=10000)
+figure2 = plot_source_energy(this=my_dt_source, figure=figure1, n_samples=10000)
 
 figure2.show()

src/openmc_source_plotter/core.py

@@ -22,16 +22,19 @@
     raise ImportError(msg)
 
 
-def sample_initial_particles(self, n_samples: int = 1000, prn_seed: int = None):
+def sample_initial_particles(this, n_samples: int = 1000, prn_seed: int = None):
+    """smaples particles from the source.
 
+    Args:
+        this: The openmc source, settings or model containing the source to plot
+        n_samples: The number of source samples to obtain.
+        prn_seed: The pseudorandom number seed.
+    """
     with TemporaryDirectory() as tmpdir:
-
-        if isinstance(self, openmc.Model):
-
-            model = self
+        if isinstance(this, openmc.Model):
+            model = this
 
         else:
-
             model = openmc.Model()
 
             materials = openmc.Materials()
@@ -42,16 +45,14 @@ def sample_initial_particles(self, n_samples: int = 1000, prn_seed: int = None):
             geometry = openmc.Geometry([cell])
             model.geometry = geometry
 
-            if isinstance(self, openmc.Settings):
-
-                model.settings = self
+            if isinstance(this, openmc.Settings):
+                model.settings = this
 
             else:  # source object
-
                 settings = openmc.Settings()
                 settings.particles = 1
                 settings.batches = 1
-                settings.source = self
+                settings.source = this
                 model.settings = settings
 
         model.export_to_model_xml()
@@ -66,7 +67,7 @@ def sample_initial_particles(self, n_samples: int = 1000, prn_seed: int = None):
 
 
 def plot_source_energy(
-    self,
+    this,
     figure: plotly.graph_objects.Figure = None,
     n_samples: int = 2000,
     prn_seed: int = 1,
@@ -79,6 +80,7 @@ def plot_source_energy(
     """makes a plot of the initial creation positions of an OpenMC source
 
     Args:
+        this: The openmc source, settings or model containing the source to plot
         figure: Optional base plotly figure to use for the plot. Passing in
             a pre made figure allows one to build up plots with from
             multiple sources. Defaults to None which makes a new figure for
@@ -110,16 +112,16 @@ def plot_source_energy(
             showlegend=True,
         )
 
-    data = self.sample_initial_particles(n_samples, prn_seed)
+    data = sample_initial_particles(this, n_samples, prn_seed)
 
     e_values = [particle.E for particle in data]
 
     # Calculate pdf for source energies
     probability, bin_edges = np.histogram(e_values, bins=energy_bins, density=True)
 
     # scaling by strength
-    if isinstance(self, openmc.SourceBase):
-        probability = probability * self.strength
+    if isinstance(this, openmc.SourceBase):
+        probability = probability * this.strength
     energy = bin_edges[:-1]
     if xaxis_units == "MeV":
         energy = energy / 1e6
@@ -138,14 +140,15 @@ def plot_source_energy(
 
 
 def plot_source_position(
-    self,
+    this: typing.Union[openmc.SourceBase, openmc.Settings, openmc.Model],
     figure=None,
     n_samples: int = 2000,
     prn_seed: int = 1,
 ):
     """makes a plot of the initial creation positions of an OpenMC source(s)
 
     Args:
+        this: The openmc source, settings or model containing the source to plot
         figure: Optional base plotly figure to use for the plot. Passing in
             a pre made figure allows one to build up plots with from
             multiple sources. Defaults to None which makes a new figure for
@@ -163,7 +166,7 @@ def plot_source_position(
             showlegend=True,
         )
 
-    data = self.sample_initial_particles(n_samples, prn_seed)
+    data = sample_initial_particles(this, n_samples, prn_seed)
 
     text = ["Energy = " + str(particle.E) + " eV" for particle in data]
 
@@ -188,14 +191,15 @@ def plot_source_position(
 
 
 def plot_source_direction(
-    self,
+    this: typing.Union[openmc.SourceBase, openmc.Settings, openmc.Model],
     figure=None,
     n_samples: int = 2000,
     prn_seed: int = 1,
 ):
     """makes a plot of the initial creation positions of an OpenMC source(s)
 
     Args:
+        this: The openmc source, settings or model containing the source to plot
         figure: Optional base plotly figure to use for the plot. Passing in
             a pre made figure allows one to build up plots with from
             multiple sources. Defaults to None which makes a new figure for
@@ -209,7 +213,7 @@ def plot_source_direction(
     figure = plotly.graph_objects.Figure()
     figure.update_layout(title="Particle initial directions")
 
-    data = self.sample_initial_particles(n_samples, prn_seed)
+    data = sample_initial_particles(this, n_samples, prn_seed)
 
     biggest_coord = max(
         max([particle.r[0] for particle in data]),
@@ -253,30 +257,3 @@ def plot_source_direction(
     )
 
     return figure
-
-
-"""
-Extents the openmc.Source class to add source plotting
-methods for energy, direction and position. Source sampling methods are
-also provided for convenience. Additional methods are plot_source_energy(),
-plot_source_position(), plot_source_direction(), sample_initial_particles()
-"""
-openmc.SourceBase.sample_initial_particles = sample_initial_particles
-openmc.model.Model.sample_initial_particles = sample_initial_particles
-openmc.Model.sample_initial_particles = sample_initial_particles
-openmc.Settings.sample_initial_particles = sample_initial_particles
-
-openmc.SourceBase.plot_source_energy = plot_source_energy
-openmc.model.Model.plot_source_energy = plot_source_energy
-openmc.Model.plot_source_energy = plot_source_energy
-openmc.Settings.plot_source_energy = plot_source_energy
-
-openmc.SourceBase.plot_source_position = plot_source_position
-openmc.model.Model.plot_source_position = plot_source_position
-openmc.Model.plot_source_position = plot_source_position
-openmc.Settings.plot_source_position = plot_source_position
-
-openmc.SourceBase.plot_source_direction = plot_source_direction
-openmc.model.Model.plot_source_direction = plot_source_direction
-openmc.Model.plot_source_direction = plot_source_direction
-openmc.Settings.plot_source_direction = plot_source_direction