functions instead of extending

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,9 @@
 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)
 
 # 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
+import 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()

src/openmc_source_plotter/core.py

@@ -22,13 +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):
+        if isinstance(this, openmc.Model):
 
-            model = self
+            model = this
 
         else:
 
@@ -42,16 +48,16 @@ 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):
+            if isinstance(this, openmc.Settings):
 
-                model.settings = self
+                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 +72,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 +85,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 +117,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 +145,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 +171,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 +196,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 +218,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 +262,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

tests/test_core_with_model.py

@@ -1,5 +1,5 @@
 import openmc
-import openmc_source_plotter
+from openmc_source_plotter import sample_initial_particles, plot_source_energy, plot_source_position
 import numpy as np
 import plotly.graph_objects as go
 import pytest
@@ -8,7 +8,7 @@
 @pytest.fixture
 def test_model():
     # initialises a new source object
-    my_source = openmc.Source()
+    my_source = openmc.IndependentSource()
 
     # sets the location of the source to x=0 y=0 z=0
     my_source.space = openmc.stats.Point((1.0, 2.0, 3.0))
@@ -38,7 +38,7 @@ def test_model():
 
 
 def test_sample_initial_particles(test_model):
-    particles = test_model.sample_initial_particles(n_samples=43)
+    particles = sample_initial_particles(this=test_model, n_samples=43)
     for particle in particles:
         assert particle.E == 15e6
         assert str(particle.particle) == "photon"
@@ -47,43 +47,44 @@ def test_sample_initial_particles(test_model):
 
 
 def test_energy_plot_with_bins(test_model):
-    plot = test_model.plot_source_energy(
+    plot = plot_source_energy(
+        this=test_model,
         n_samples=10,
         energy_bins=np.linspace(0, 20e6, 100),
     )
     assert isinstance(plot, go.Figure)
 
 
 def test_energy_plot(test_model):
-    plot = test_model.plot_source_energy(n_samples=10)
+    plot = plot_source_energy(this=test_model, n_samples=10)
     assert isinstance(plot, go.Figure)
     assert len(plot.data[0]["x"]) == 1
 
 
 def test_position_plot(test_model):
-    plot = test_model.plot_source_position(n_samples=10)
+    plot = plot_source_position(this=test_model, n_samples=10)
     assert isinstance(plot, go.Figure)
 
 
 def test_direction_plot(test_model):
-    plot = test_model.plot_source_direction(n_samples=10)
+    plot = plot_source_direction(this=test_model, n_samples=10)
     assert isinstance(plot, go.Figure)
 
 
 def test_energy_plot_with_figure(test_model):
     base_figure = go.Figure()
-    plot = test_model.plot_source_energy(figure=base_figure, n_samples=10)
+    plot = plot_source_energy(this=test_model, figure=base_figure, n_samples=10)
     assert isinstance(plot, go.Figure)
     assert len(plot.data[0]["x"]) == 1
 
 
 def test_position_plot_with_figure(test_model):
     base_figure = go.Figure()
-    plot = test_model.plot_source_position(figure=base_figure, n_samples=10)
+    plot = plot_source_position(this=test_model, figure=base_figure, n_samples=10)
     assert isinstance(plot, go.Figure)
 
 
 def test_direction_plot_with_figure(test_model):
     base_figure = go.Figure()
-    plot = test_model.plot_source_direction(figure=base_figure, n_samples=10)
+    plot = plot_source_direction(this=test_model, figure=base_figure, n_samples=10)
     assert isinstance(plot, go.Figure)