Merge pull request #11 from fusion-energy/using_source_sample_for_energy

Using source sample for energy

README.md

@@ -8,8 +8,10 @@ pip install openmc_source_plotter
 ```
 
 temporary fix
-For fixed source sources it is currently necessary to use openmc version 0.11
-and also to point the ```openmc_exec``` path to the openmc executable
+For position and direction plotting with fixed source sources it is currently
+necessary to use openmc version 0.11 and also to point the ```openmc_exec```
+path to the openmc executable.
+
 This can be installed with:
 ```bash
 conda install -c conda-forge openmc=0.11
@@ -47,7 +49,6 @@ plot = osp.plot_source_energy(
     source=my_source,
     number_of_particles=2000,
     energy_bins=np.linspace(0, 20e6, 100),
-    openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
 )
 
 plot.show()
@@ -78,7 +79,6 @@ plot = osp.plot_source_energy(
     source=[my_dt_source, my_dd_source],
     number_of_particles=10000,
     energy_bins=np.linspace(0, 20e6, 100),
-    openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
 )
 
 plot.show()

examples/example_plot_source_energy.py

@@ -13,7 +13,6 @@
     source=my_source,
     number_of_particles=10000,
     energy_bins=np.linspace(0, 20e6, 100),
-    openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
 )
 
 plot.show()

examples/example_plot_two_source_energies.py

@@ -18,7 +18,6 @@
     source=[my_dt_source, my_dd_source],
     number_of_particles=10000,
     energy_bins=np.linspace(0, 20e6, 100),
-    openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
 )
 
 plot.show()

openmc_source_plotter/core.py

@@ -15,8 +15,7 @@
 def plot_source_energy(
     source: Union[openmc.Source, List[openmc.Source]],
     number_of_particles: int = 2000,
-    openmc_exec="openmc",
-    energy_bins: np.array = np.linspace(0, 20e6, 50),
+    energy_bins: Union[np.array, str] = "auto",
     filename: str = None,
 ):
     """makes a plot of the energy distribution OpenMC source(s)
@@ -25,10 +24,10 @@ def plot_source_energy(
         source: The openmc.Source object or list of openmc.Source objects to plot.
         number_of_particles: The number of source samples to obtain, more will
             take longer but produce a smoother plot.
-        energy_bins: A numpy array of energy bins to use as energy bins.
-        openmc_exec: The path of the openmc executable to use
+        energy_bins: A numpy array of energy bins to use as energy bins. 'Auto'
+            is also accepted and this picks the bins for you using numpy
         filename: the filename to save the plot as should end with the correct
-            extention supported by matplotlib (e.g .png) or plotly (e.g .html)
+            extension supported by matplotlib (e.g .png) or plotly (e.g .html)
     """
 
     figure = go.Figure()
@@ -37,21 +36,11 @@ def plot_source_energy(
         source = [source]
 
     for single_source in source:
-        tmp_filename = tempfile.mkstemp(suffix=".h5", prefix=f"openmc_source_")[1]
-        create_initial_particles(
-            source=single_source,
-            number_of_particles=number_of_particles,
-            openmc_exec=openmc_exec,
-            output_source_filename=tmp_filename,
-        )
-
-        print("getting particle data", tmp_filename)
-        data = get_particle_data(tmp_filename)
 
-        e_values = data["e_values"]
+        e_values = single_source.energy.sample(n_samples=number_of_particles)
 
         # Calculate pdf for source energies
-        probability, bin_edges = np.histogram(e_values, energy_bins, density=True)
+        probability, bin_edges = np.histogram(e_values, bins=energy_bins, density=True)
 
         # Plot source energy histogram
         figure.add_trace(
@@ -60,7 +49,6 @@ def plot_source_energy(
                 y=probability * np.diff(energy_bins),
                 line={"shape": "hv"},
                 hoverinfo="text",
-                name=tmp_filename,
             )
         )
 

tests/test_core.py

@@ -44,7 +44,6 @@ def test_energy_plot(self):
             source=self.my_source,
             number_of_particles=10000,
             energy_bins=np.linspace(0, 20e6, 100),
-            openmc_exec=self.openmc_exec_dict[self.current_computer],
         )
         assert isinstance(plot, go.Figure)