Added distribute packets which can replace decays_per_isotope and fra…

…ctional energy. Much Cleaner

tardis/energy_input/gamma_ray_transport.py

@@ -580,6 +580,41 @@ def calculate_energy_per_mass(decay_energy, raw_isotope_abundance, cell_masses):
     return energy_per_mass, energy_df
 
 
+def distribute_packets(decay_energy, energy_per_packet):
+
+    packets_per_isotope = {}
+    for shell, isotopes in decay_energy.items():
+        packets_per_isotope[shell] = {}
+        for name, isotope in isotopes.items():
+            packets_per_isotope[shell][name] = {}
+            for line, energy in isotope.items():
+                packets_per_isotope[shell][name][line] = int(
+                    energy / energy_per_packet
+                )
+
+    packets_per_isotope_list = []
+    for shell, parent_isotope in packets_per_isotope.items():
+        for isotopes, isotope_dict in parent_isotope.items():
+            for name, value in isotope_dict.items():
+                packets_per_isotope_list.append(
+                    {
+                        "shell": shell,
+                        "element": name,
+                        "value": value,
+                    }
+                )
+
+    df = pd.DataFrame(packets_per_isotope_list)
+    packets_per_isotope_df = pd.pivot_table(
+        df,
+        values="value",
+        index="element",
+        columns="shell",
+    )
+
+    return packets_per_isotope_df
+
+
 def packets_per_isotope(fractional_decay_energy, decayed_packet_count_dict):
     packets_per_isotope = {
         shell: {

tardis/energy_input/main_gamma_ray_loop.py

@@ -23,6 +23,7 @@
     get_taus,
     fractional_decay_energy,
     packets_per_isotope,
+    distribute_packets,
 )
 
 logger = logging.getLogger(__name__)
@@ -131,17 +132,26 @@ def run_gamma_ray_loop(
         total_isotope_number, axis=0
     )
 
-    decayed_packet_count_dict = decayed_packet_count.to_dict()
-    fractional_decay_energy_dict = fractional_decay_energy(decayed_energy)
+    # decayed_packet_count_dict = decayed_packet_count.to_dict()
+    # fractional_decay_energy_dict = fractional_decay_energy(decayed_energy)
+    # packets_per_isotope_df = (
+    #    packets_per_isotope(
+    #        fractional_decay_energy_dict, decayed_packet_count_dict
+    #    )
+    #    .round()
+    #    .fillna(0)
+    #    .astype(int)
+    # )
+
+    total_energy = energy_df.sum().sum()
+    energy_per_packet = total_energy / num_decays
     packets_per_isotope_df = (
-        packets_per_isotope(
-            fractional_decay_energy_dict, decayed_packet_count_dict
-        )
+        distribute_packets(decayed_energy, energy_per_packet)
         .round()
         .fillna(0)
         .astype(int)
     )
-    total_energy = energy_df.sum().sum()
+
     total_energy = total_energy * u.eV.to("erg")
 
     logger.info(f"Total gamma-ray energy is {total_energy}")