diff --git a/src/openmc_source_plotter/core.py b/src/openmc_source_plotter/core.py index 40dbb20..34e7f83 100644 --- a/src/openmc_source_plotter/core.py +++ b/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 diff --git a/tests/test_core_with_settings.py b/tests/test_core_with_settings.py new file mode 100644 index 0000000..59d83ed --- /dev/null +++ b/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)