Merge pull request #37 from fusion-energy/adding_support_for_settings…

….plot

added plotting methods to openmc.Settings

src/openmc_source_plotter/core.py

@@ -3,7 +3,7 @@
 """Provides functions for plotting source information"""
 
 import typing
-
+from tempfile import TemporaryDirectory
 import numpy as np
 import openmc
 import openmc.lib
@@ -24,34 +24,43 @@
 
 def sample_initial_particles(self, n_samples: int = 1000, prn_seed: int = None):
 
-    if isinstance(self, openmc.Model):
+    with TemporaryDirectory() as tmpdir:
 
-        self.settings.export_to_xml()
-        self.materials.export_to_xml()
-        self.geometry.export_to_xml()
+        if isinstance(self, openmc.Model):
 
-    else:  # source object
+            model = self
 
-        settings = openmc.Settings()
-        settings.particles = 1
-        settings.batches = 1
-        settings.source = self
+        else:
 
-        materials = openmc.Materials()
+            model = openmc.Model()
 
-        sph = openmc.Sphere(r=9999999999, boundary_type="vacuum")
-        cell = openmc.Cell(region=-sph)
-        geometry = openmc.Geometry([cell])
+            materials = openmc.Materials()
+            model.materials = materials
 
-        settings.export_to_xml()
-        materials.export_to_xml()
-        geometry.export_to_xml()
+            sph = openmc.Sphere(r=9999999999, boundary_type="vacuum")
+            cell = openmc.Cell(region=-sph)
+            geometry = openmc.Geometry([cell])
+            model.geometry = geometry
 
-    openmc.lib.init(output=False)
-    particles = openmc.lib.sample_external_source(
-        n_samples=n_samples, prn_seed=prn_seed
-    )
-    openmc.lib.finalize()
+            if isinstance(self, openmc.Settings):
+
+                model.settings = self
+
+            else:  # source object
+
+                settings = openmc.Settings()
+                settings.particles = 1
+                settings.batches = 1
+                settings.source = self
+                model.settings = settings
+
+        model.export_to_model_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
 
@@ -242,15 +251,19 @@ def plot_source_direction(
 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_settings.py

@@ -0,0 +1,81 @@
+import openmc
+import openmc_source_plotter
+import numpy as np
+import plotly.graph_objects as go
+import pytest
+
+
+@pytest.fixture
+def test_settings():
+    # 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
+
+    return settings
+
+
+def test_sample_initial_particles(test_settings):
+    particles = test_settings.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_settings):
+    plot = test_settings.plot_source_energy(
+        n_samples=10,
+        energy_bins=np.linspace(0, 20e6, 100),
+    )
+    assert isinstance(plot, go.Figure)
+
+
+def test_energy_plot(test_settings):
+    plot = test_settings.plot_source_energy(n_samples=10)
+    assert isinstance(plot, go.Figure)
+    assert len(plot.data[0]["x"]) == 1
+
+
+def test_position_plot(test_settings):
+    plot = test_settings.plot_source_position(n_samples=10)
+    assert isinstance(plot, go.Figure)
+
+
+def test_direction_plot(test_settings):
+    plot = test_settings.plot_source_direction(n_samples=10)
+    assert isinstance(plot, go.Figure)
+
+
+def test_energy_plot_with_figure(test_settings):
+    base_figure = go.Figure()
+    plot = test_settings.plot_source_energy(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_settings):
+    base_figure = go.Figure()
+    plot = test_settings.plot_source_position(figure=base_figure, n_samples=10)
+    assert isinstance(plot, go.Figure)
+
+
+def test_direction_plot_with_figure(test_settings):
+    base_figure = go.Figure()
+    plot = test_settings.plot_source_direction(figure=base_figure, n_samples=10)
+    assert isinstance(plot, go.Figure)