From fa458859de24b0f4ff2985518f11cf05b07adacf Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Fri, 10 May 2024 15:14:46 +0100
Subject: [PATCH 1/2] special case for mixed stats

---
 src/openmc_source_plotter/material.py | 51 ++++++++++++++++-----------
 1 file changed, 31 insertions(+), 20 deletions(-)

diff --git a/src/openmc_source_plotter/material.py b/src/openmc_source_plotter/material.py
index ab9df53..a8eb595 100644
--- a/src/openmc_source_plotter/material.py
+++ b/src/openmc_source_plotter/material.py
@@ -1,6 +1,32 @@
 import openmc
 import matplotlib.pyplot as plt
 
+def _get_energy_prob(energy_dis):
+    """gets the energy and probability for different openmc.stats including the
+     openmc.stats.Mixutre which itself is made from openmc.stats"""
+
+    if isinstance(energy_dis, openmc.stats.Mixture):
+        stats = energy_dis.distribution
+        multipliers = energy_dis.probability
+    else:
+        stats = [energy_dis]
+        multipliers = [1.]
+
+    probs = []
+    en = []
+
+    for stat, multiplier in zip(stats, multipliers):
+
+        for p in stat.p:
+            probs.append(0)
+            probs.append(p*multiplier)
+            probs.append(0)
+        for x in stat.x:
+            en.append(x)
+            en.append(x)
+            en.append(x)
+
+    return en, probs
 
 def plot_gamma_emission(
     material,
@@ -48,16 +74,9 @@ def plot_gamma_emission(
         probs = []
         en = []
         energy_dis = material.get_decay_photon_energy(clip_tolerance=0.0)
-        for p in energy_dis.p:
-            probs.append(0)
-            probs.append(p)
-            probs.append(0)
-        for x in energy_dis.x:
-            en.append(x)
-            en.append(x)
-            en.append(x)
-        # print(en)
-        # print(probs)
+
+        en, probs = _get_energy_prob(energy_dis)
+
         lineid_plot.plot_line_ids(
             en,
             # material.decay_photon_energy.x,
@@ -68,17 +87,9 @@ def plot_gamma_emission(
         )
 
     else:
-        probs = []
-        en = []
         energy_dis = material.get_decay_photon_energy(clip_tolerance=0.0)
-        for p in energy_dis.p:
-            probs.append(0)
-            probs.append(p)
-            probs.append(0)
-        for x in energy_dis.x:
-            en.append(x)
-            en.append(x)
-            en.append(x)
+
+        en, probs = _get_energy_prob(energy_dis)
 
         # plt.scatter(energy_dis.x, energy_dis.p)
         plt.plot(en, probs)

From 06f507666e2ba9e4d92b0a54566bedacb35e427f Mon Sep 17 00:00:00 2001
From: shimwell <shimwell@users.noreply.github.com>
Date: Fri, 10 May 2024 14:15:06 +0000
Subject: [PATCH 2/2] [skip ci] Apply formatting changes

---
 src/openmc_source_plotter/material.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/openmc_source_plotter/material.py b/src/openmc_source_plotter/material.py
index a8eb595..bfdf273 100644
--- a/src/openmc_source_plotter/material.py
+++ b/src/openmc_source_plotter/material.py
@@ -1,16 +1,17 @@
 import openmc
 import matplotlib.pyplot as plt
 
+
 def _get_energy_prob(energy_dis):
     """gets the energy and probability for different openmc.stats including the
-     openmc.stats.Mixutre which itself is made from openmc.stats"""
+    openmc.stats.Mixutre which itself is made from openmc.stats"""
 
     if isinstance(energy_dis, openmc.stats.Mixture):
         stats = energy_dis.distribution
         multipliers = energy_dis.probability
     else:
         stats = [energy_dis]
-        multipliers = [1.]
+        multipliers = [1.0]
 
     probs = []
     en = []
@@ -19,7 +20,7 @@ def _get_energy_prob(energy_dis):
 
         for p in stat.p:
             probs.append(0)
-            probs.append(p*multiplier)
+            probs.append(p * multiplier)
             probs.append(0)
         for x in stat.x:
             en.append(x)
@@ -28,6 +29,7 @@ def _get_energy_prob(energy_dis):
 
     return en, probs
 
+
 def plot_gamma_emission(
     material,
     label_top: int = None,