From 2704ee7927c1e36d128ad8ec8dfd89e0225192df Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Sat, 9 Mar 2024 12:17:08 -0500
Subject: [PATCH] Added distribute packets which can replace decays_per_isotope
 and fractional energy. Much Cleaner

---
 tardis/energy_input/gamma_ray_transport.py | 35 ++++++++++++++++++++++
 tardis/energy_input/main_gamma_ray_loop.py | 22 ++++++++++----
 2 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/tardis/energy_input/gamma_ray_transport.py b/tardis/energy_input/gamma_ray_transport.py
index 7f2243d451d..2da2d1e21e5 100644
--- a/tardis/energy_input/gamma_ray_transport.py
+++ b/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: {
diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index 0530b6abe5b..eaeec4e2cc3 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/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}")