Merge pull request #32 from fusion-energy/adding_sample_source_to_model_class

adds sample_initial_particles to the openmc.Model method

Author: shimwell

README.md

@@ -12,14 +12,16 @@ pip install openmc_source_plotter
 
 # Features
 
-The package simply extends the default ```openmc.Source``` to provides additional functions that:
+The package simply extends the default ```openmc.IndependentSourceBase``` to provides additional functions that:
 
 - extract the positions, directions and energy of particles
-- visualise an ```osp.SourceWithPlotting``` with respect to:
+- visualise an ```openmc.IndependentSourceBase``` with respect to:
   - direction
   - energy
   - position
 
+Or just sample the with ```openmc.Model.sample_initial_particles```
+
 # Example plots
 
 Below are some basic examples, for more examples see the [examples folder](https://github.com/fusion-energy/openmc_source_plotter/tree/main/examples) for example usage scripts.
@@ -29,13 +31,13 @@ Below are some basic examples, for more examples see the [examples folder](https
 
 ```python
 import openmc
-import openmc_source_plotter  # extends openmc.Source with plotting functions
+import openmc_source_plotter  # extends openmc.IndependentSource with plotting functions
 
 # initialises a new source object
-my_source = openmc.Source()
+my_source = openmc.IndependentSource()
 
-# sets the energy distribution to a Muir distribution neutrons for DT fusion 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 for DT fusion 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=2000)
@@ -49,18 +51,18 @@ plot.show()
 
 ```python
 import openmc
-import openmc_source_plotter  # extends openmc.Source with plotting functions
+import openmc_source_plotter  # extends openmc.IndependentSource with plotting functions
 
 # initialises a new source object
-my_dt_source = openmc.Source()
+my_dt_source = openmc.IndependentSource()
 
-# 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)
+my_dd_source = openmc.IndependentSource()
+# 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 by building on the first plot
 figure1 = my_dd_source.plot_source_energy(n_samples=10000)
@@ -75,10 +77,10 @@ figure2.show()
 
 ```python
 import openmc
-import openmc_source_plotter  # extends openmc.Source with plotting functions
+import openmc_source_plotter  # extends openmc.IndependentSource with plotting functions
 
 # initializes a new source object
-my_source = openmc.Source()
+my_source = openmc.IndependentSource()
 
 # sets the direction to isotropic
 my_source.angle = openmc.stats.Isotropic()
@@ -92,7 +94,7 @@ plot.show()
 ![openmc particle source direction plot](https://user-images.githubusercontent.com/8583900/143615706-3b3a8467-0233-42d6-a66c-d536c80a01d8.png)
 
 
-## Plot gamma spectrum of particles
+<!-- ## Plot gamma spectrum of particles
 
 ```python
 import openmc
@@ -114,16 +116,16 @@ plt.xscale("log")  # modify axis from default settings
 plt.savefig("gamma_spec.png")
 ```
 
-![openmc gamma spectrum](https://user-images.githubusercontent.com/8583900/228280129-b8160e18-9ca9-4b20-a4e1-d2948908daf6.png)
+![openmc gamma spectrum](https://user-images.githubusercontent.com/8583900/228280129-b8160e18-9ca9-4b20-a4e1-d2948908daf6.png) -->
 
 ## Plot position of particles
 
 ```python
 import openmc
-import openmc_source_plotter  # extends openmc.Source with plotting functions
+import openmc_source_plotter  # extends openmc.IndependentSource with plotting functions
 
 # initialises a new source object
-my_source = openmc.Source()
+my_source = openmc.IndependentSource()
 
 # the distribution of radius is just a single value
 radius = openmc.stats.Discrete([10], [1])
@@ -149,6 +151,39 @@ plot.show()
 
 ![openmc particle source position plot](https://user-images.githubusercontent.com/8583900/179424915-bee56a87-6214-46ef-8625-92b8f4cbd1b3.png)
 
+
+## Extract particle objects
+
+A list of ```openmc.Particle``` objects can be obtained using ```model.sample_initial_particles()``` or ```openmc.SourceBase.sample_initial_particles()```
+
+```python
+import openmc
+import openmc_source_plotter  # extents openmc.Model with sample_initial_particles method
+
+settings = openmc.Settings()
+settings.particles = 1
+settings.batches = 1
+my_source = openmc.IndependentSource()
+my_source.energy = openmc.stats.muir(e0=14080000.0, m_rat=5.0, kt=20000.0)
+settings.source = my_source
+materials = openmc.Materials()
+sph = openmc.Sphere(r=100, boundary_type="vacuum")
+cell = openmc.Cell(region=-sph)
+geometry = openmc.Geometry([cell])
+
+model = openmc.Model(geometry, materials, settings)
+
+particles = model.sample_initial_particles(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>]
+
+print(particles[0].E)
+>>>1.440285e+07
+```
+
+## Related packages
+
 Tokamak sources can also be plotted using the [openmc-plasma-source](https://github.com/fusion-energy/openmc-plasma-source) package
 
 ![openmc_source_plotter_openmc-plasma-source_tokamak](https://user-images.githubusercontent.com/8583900/187487894-ba0bd025-46f2-4c7d-8b15-3d260aed47a0.png)

src/openmc_source_plotter/core.py

@@ -11,28 +11,45 @@
 
 
 def sample_initial_particles(self, n_samples: int = 1000, prn_seed: int = None):
-    settings = openmc.Settings()
-    settings.particles = 1
-    settings.batches = 1
-    settings.source = self
-    settings.export_to_xml()
 
-    materials = openmc.Materials()
-    materials.export_to_xml()
+    if isinstance(self, openmc.Model):
 
-    sph = openmc.Sphere(r=9999999999, boundary_type="vacuum")
-    cell = openmc.Cell(region=-sph)
-    geometry = openmc.Geometry([cell])
+        self.settings.export_to_xml()
+        self.materials.export_to_xml()
+        self.geometry.export_to_xml()
 
-    geometry.export_to_xml()
+        openmc.lib.init(output=False)
+        particles = openmc.lib.sample_external_source(
+            n_samples=n_samples, prn_seed=prn_seed
+        )
+        openmc.lib.finalize()
 
-    openmc.lib.init(output=False)
-    particles = openmc.lib.sample_external_source(
-        n_samples=n_samples, prn_seed=prn_seed
-    )
-    openmc.lib.finalize()
+        return particles
+
+    else:  # source object
+
+        settings = openmc.Settings()
+        settings.particles = 1
+        settings.batches = 1
+        settings.source = self
+        settings.export_to_xml()
+
+        materials = openmc.Materials()
+        materials.export_to_xml()
 
-    return particles
+        sph = openmc.Sphere(r=9999999999, boundary_type="vacuum")
+        cell = openmc.Cell(region=-sph)
+        geometry = openmc.Geometry([cell])
+
+        geometry.export_to_xml()
+
+        openmc.lib.init(output=False)
+        particles = openmc.lib.sample_external_source(
+            n_samples=n_samples, prn_seed=prn_seed
+        )
+        openmc.lib.finalize()
+
+        return particles
 
 
 def plot_source_energy(
@@ -218,6 +235,9 @@ def plot_source_direction(
 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.SourceBase.plot_source_energy = plot_source_energy
 openmc.SourceBase.plot_source_position = plot_source_position
 openmc.SourceBase.plot_source_direction = plot_source_direction

tests/test_core.py

@@ -11,21 +11,68 @@ def test_source():
     my_source = openmc.Source()
 
     # sets the location of the source to x=0 y=0 z=0
-    my_source.space = openmc.stats.Point((0, 0, 0))
+    my_source.space = openmc.stats.Point((4.0, 5.0, 6.0))
 
     # sets the direction to isotropic
     my_source.angle = openmc.stats.Isotropic()
 
     # sets the energy distribution to 100% 14MeV neutrons
     my_source.energy = openmc.stats.Discrete([14e6], [1])
 
+    my_source.particle = "neutron"
+
     my_source = my_source
     return my_source
 
 
-def test_sample_initial_particles(test_source):
-    data = test_source.sample_initial_particles(n_samples=42)
-    assert len(data) == 42
+@pytest.fixture
+def test_model():
+    # initialises a new source object
+    my_source = openmc.Source()
+
+    # 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))
+
+    # sets the direction to isotropic
+    my_source.angle = openmc.stats.Isotropic()
+
+    # sets the energy distribution to 100% 14MeV neutrons
+    my_source.energy = openmc.stats.Discrete([15e6], [1])
+
+    my_source.particle = "photon"
+
+    settings = openmc.Settings()
+    settings.particles = 1
+    settings.batches = 1
+    settings.source = my_source
+
+    materials = openmc.Materials()
+
+    sph = openmc.Sphere(r=9999999999, boundary_type="vacuum")
+    cell = openmc.Cell(region=-sph)
+    geometry = openmc.Geometry([cell])
+
+    model = openmc.Model(geometry, materials, settings)
+
+    return model
+
+
+def test_source_sample_initial_particles(test_source):
+    particles = test_source.sample_initial_particles(n_samples=42)
+    for particle in particles:
+        assert particle.E == 14e6
+        assert str(particle.particle) == "neutron"
+        assert particle.r == (4.0, 5.0, 6.0)
+    assert len(particles) == 42
+
+
+def test_model_sample_initial_particles(test_model):
+    particles = test_model.sample_initial_particles(n_samples=43)
+    for particle in particles:
+        assert particle.E == 15e6
+        assert str(particle.particle) == "photon"
+        assert particle.r == (1.0, 2.0, 3.0)
+    assert len(particles) == 43
 
 
 def test_energy_plot_with_bins(test_source):