From 56d8f44620977521d5f2ae8908c50c7cd63d0c94 Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Wed, 24 Apr 2024 16:16:15 -0400
Subject: [PATCH 01/18] Added decay of abundances with time

---
 tardis/energy_input/gamma_ray_channel.py | 36 ++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/tardis/energy_input/gamma_ray_channel.py b/tardis/energy_input/gamma_ray_channel.py
index 6bbb7a9685a..40c595f3ada 100644
--- a/tardis/energy_input/gamma_ray_channel.py
+++ b/tardis/energy_input/gamma_ray_channel.py
@@ -5,6 +5,7 @@
 import radioactivedecay as rd
 
 from tardis.energy_input.util import KEV2ERG
+from tardis.model.matter.decay import IsotopicMassFraction
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
@@ -167,3 +168,38 @@ def create_isotope_decay_df(cumulative_decay_df, gamma_ray_lines):
     )
 
     return isotope_decay_df
+
+
+def evolve_mass_fraction(raw_isotope_mass_fraction, time_array):
+    """
+    Function to evolve the mass fraction of isotopes with time.
+
+    Parameters
+    ----------
+    raw_isotope_mass_fraction : pd.DataFrame
+        isotope mass fraction in mass fractions.
+    time_array : np.array
+        array of time in days.
+
+    Returns
+    -------
+    time_evolved_isotope_mass_fraction : pd.DataFrame
+        time evolved mass fraction of isotopes.
+    """
+
+    initial_isotope_mass_fraction = raw_isotope_mass_fraction
+    isotope_mass_fraction_list = []
+
+    for time in time_array:
+
+        decayed_isotope_mass_fraction = IsotopicMassFraction(
+            initial_isotope_mass_fraction
+        ).decay(time)
+        isotope_mass_fraction_list.append(decayed_isotope_mass_fraction)
+        initial_isotope_mass_fraction = decayed_isotope_mass_fraction
+
+    time_evolved_isotope_mass_fraction = pd.concat(
+        isotope_mass_fraction_list, keys=time_array, names=["time"]
+    )
+
+    return time_evolved_isotope_mass_fraction

From ddb44a24347c9fee4c42a6a479fa6dd69523ee05 Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Fri, 3 May 2024 11:20:20 -0400
Subject: [PATCH 02/18] Adding dataframe to gammaraypacketsource

---
 tardis/energy_input/gamma_packet_loop.py      |  71 +++----
 tardis/energy_input/gamma_ray_channel.py      |  57 +++++
 .../energy_input/gamma_ray_packet_source.py   | 191 +++++++++++------
 tardis/energy_input/gamma_ray_transport.py    |   2 +-
 tardis/energy_input/main_gamma_ray_loop.py    | 198 +++++++++++++++++-
 tardis/energy_input/util.py                   |  45 ++++
 tardis/model/matter/decay.py                  |   4 +-
 7 files changed, 465 insertions(+), 103 deletions(-)

diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py
index 8f76d1f2c07..17f93775d44 100644
--- a/tardis/energy_input/gamma_packet_loop.py
+++ b/tardis/energy_input/gamma_packet_loop.py
@@ -1,3 +1,4 @@
+import logging
 import numpy as np
 from numba import njit
 
@@ -30,6 +31,9 @@
 )
 from tardis.energy_input.gamma_ray_estimators import deposition_estimator_kasen
 
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
 
 @njit(**njit_dict_no_parallel)
 def gamma_packet_loop(
@@ -47,8 +51,6 @@ def gamma_packet_loop(
     dt_array,
     effective_time_array,
     energy_bins,
-    energy_df_rows,
-    energy_plot_df_rows,
     energy_out,
     packets_info_array,
 ):
@@ -104,13 +106,15 @@ def gamma_packet_loop(
     escaped_packets = 0
     scattered_packets = 0
     packet_count = len(packets)
-    print("Entering gamma ray loop for " + str(packet_count) + " packets")
 
-    deposition_estimator = np.zeros_like(energy_df_rows)
+    print("Packet count: ", packet_count)
+    # deposition_estimator = np.zeros_like(energy_df_rows)
 
     for i in range(packet_count):
         packet = packets[i]
         time_index = get_index(packet.time_current, times)
+        # print("Time index: ", time_index)
+        # print("Time current: ", packet.time_current)
 
         if time_index < 0:
             print(packet.time_current, time_index)
@@ -215,16 +219,16 @@ def gamma_packet_loop(
 
             packet = move_packet(packet, distance)
 
-            deposition_estimator[packet.shell, time_index] += (
-                (initial_energy * 1000)
-                * distance
-                * (packet.energy_cmf / initial_energy)
-                * deposition_estimator_kasen(
-                    comoving_energy,
-                    mass_density_time[packet.shell, time_index],
-                    iron_group_fraction_per_shell[packet.shell],
-                )
-            )
+            # deposition_estimator[packet.shell, time_index] += (
+            #    (initial_energy * 1000)
+            #    * distance
+            #    * (packet.energy_cmf / initial_energy)
+            #    * deposition_estimator_kasen(
+            #        comoving_energy,
+            #       mass_density_time[packet.shell, time_index],
+            #        iron_group_fraction_per_shell[packet.shell],
+            #    )
+            # )
 
             if distance == distance_time:
                 time_index += 1
@@ -249,24 +253,24 @@ def gamma_packet_loop(
 
                 # Save packets to dataframe rows
                 # convert KeV to eV / s / cm^3
-                energy_df_rows[packet.shell, time_index] += (
-                    ejecta_energy_gained * 1000
-                )
-
-                energy_plot_df_rows[i] = np.array(
-                    [
-                        i,
-                        ejecta_energy_gained * 1000
-                        # * inv_volume_time[packet.shell, time_index]
-                        / dt,
-                        packet.get_location_r(),
-                        packet.time_current,
-                        packet.shell,
-                        compton_opacity,
-                        photoabsorption_opacity,
-                        pair_creation_opacity,
-                    ]
-                )
+                # energy_df_rows[packet.shell, time_index] += (
+                #    ejecta_energy_gained * 1000
+                # )
+
+                # energy_plot_df_rows[i] = np.array(
+                #    [
+                #        i,
+                #        ejecta_energy_gained * 1000
+                # * inv_volume_time[packet.shell, time_index]
+                #        / dt,
+                #        packet.get_location_r(),
+                #        packet.time_current,
+                #        packet.shell,
+                #        compton_opacity,
+                #        photoabsorption_opacity,
+                #        pair_creation_opacity,
+                #    ]
+                # )
 
                 if packet.status == GXPacketStatus.PHOTOABSORPTION:
                     # Packet destroyed, go to the next packet
@@ -314,10 +318,7 @@ def gamma_packet_loop(
     print("Scattered packets:", scattered_packets)
 
     return (
-        energy_df_rows,
-        energy_plot_df_rows,
         energy_out,
-        deposition_estimator,
         bin_width,
         packets_info_array,
     )
diff --git a/tardis/energy_input/gamma_ray_channel.py b/tardis/energy_input/gamma_ray_channel.py
index 40c595f3ada..42df619d4ed 100644
--- a/tardis/energy_input/gamma_ray_channel.py
+++ b/tardis/energy_input/gamma_ray_channel.py
@@ -203,3 +203,60 @@ def evolve_mass_fraction(raw_isotope_mass_fraction, time_array):
     )
 
     return time_evolved_isotope_mass_fraction
+
+
+def time_evolve_cumulative_decays(
+    raw_isotope_mass_fraction, shell_masses, gamma_ray_lines, time_array
+):
+    """
+    Function to calculate the total decays for each isotope for each shell at each time step.
+
+    Parameters
+    ----------
+    raw_isotope_mass_fraction : pd.DataFrame
+        isotope abundance in mass fractions.
+    shell_masses : numpy.ndarray
+        shell masses in units of g
+    gamma_ray_lines : pd.DataFrame
+        gamma ray lines from nndc stored as a pandas dataframe.
+    time_array : numpy.ndarray
+        array of time steps in days.
+
+    Returns
+    -------
+    time_evolve_decay_df : pd.DataFrame
+        dataframe of isotopes for each shell with their decay mode, number of decays, radiation type,
+        radiation energy and radiation intensity at each time step.
+
+    """
+
+    isotope_mass_fraction_list = []
+    cumulative_decay_df_list = []
+    initial_isotope_mass_fraction = raw_isotope_mass_fraction
+
+    for time in time_array:
+        decayed_isotope_mass_fraction = IsotopicMassFraction(
+            initial_isotope_mass_fraction
+        ).decay(time)
+
+        isotope_dict = create_isotope_dicts(
+            decayed_isotope_mass_fraction, shell_masses
+        )
+        inventories = create_inventories_dict(isotope_dict)
+        cumulative_decay_df = calculate_total_decays(inventories, time)
+        isotope_decay_df = create_isotope_decay_df(
+            cumulative_decay_df, gamma_ray_lines
+        )
+        isotope_mass_fraction_list.append(isotope_decay_df)
+        cumulative_decay_df_list.append(cumulative_decay_df)
+        initial_isotope_mass_fraction = decayed_isotope_mass_fraction
+
+    time_evolved_cumulative_decay = pd.concat(
+        cumulative_decay_df_list, keys=time_array, names=["time"]
+    )
+
+    time_evolve_decay_df = pd.concat(
+        isotope_mass_fraction_list, keys=time_array, names=["time"]
+    )
+
+    return time_evolve_decay_df, time_evolved_cumulative_decay
diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index a5d1522043e..1bcc008a4e1 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -1,5 +1,6 @@
 import numpy as np
 import pandas as pd
+from numba import njit
 
 from tardis.energy_input.energy_source import (
     positronium_continuum,
@@ -11,10 +12,13 @@
 from tardis.energy_input.util import (
     H_CGS_KEV,
     doppler_factor_3d,
+    doppler_factor_3D_all_packets,
     get_index,
     get_random_unit_vector,
+    C_CGS,
 )
 from tardis.montecarlo.packet_source import BasePacketSource
+from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel
 
 
 class RadioactivePacketSource(BasePacketSource):
@@ -435,18 +439,18 @@ def create_packets(self, decays_per_isotope, *args, **kwargs):
                 packet_energies_cmf[
                     packet_index - isotope_packet_count : packet_index
                 ] = initial_packet_energies_cmf
-                nus_rf[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_nus_rf
-                nus_cmf[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_nus_cmf
-                shells[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_shells
-                times[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_times
+                nus_rf[packet_index - isotope_packet_count : packet_index] = (
+                    initial_nus_rf
+                )
+                nus_cmf[packet_index - isotope_packet_count : packet_index] = (
+                    initial_nus_cmf
+                )
+                shells[packet_index - isotope_packet_count : packet_index] = (
+                    initial_shells
+                )
+                times[packet_index - isotope_packet_count : packet_index] = (
+                    initial_times
+                )
 
         return GXPacketCollection(
             locations,
@@ -487,7 +491,7 @@ class GammaRayPacketSource(BasePacketSource):
     def __init__(
         self,
         packet_energy,
-        gamma_ray_lines,
+        isotope_decay_df,
         positronium_fraction,
         inner_velocities,
         outer_velocities,
@@ -502,7 +506,7 @@ def __init__(
         **kwargs,
     ):
         self.packet_energy = packet_energy
-        self.gamma_ray_lines = gamma_ray_lines
+        self.isotope_decay_df = isotope_decay_df
         self.positronium_fraction = positronium_fraction
         self.inner_velocities = inner_velocities
         self.outer_velocities = outer_velocities
@@ -575,7 +579,7 @@ def create_packet_nus(
         zs = np.random.random(no_of_packets)
         for i in range(no_of_packets):
             # positron
-            if packets.iloc[i]["decay_type"] == "bp":
+            if packets.iloc[i]["radiation"] == "bp":
                 # positronium formation 75% of the time if fraction is 1
                 if zs[i] < positronium_fraction and np.random.random() < 0.75:
                     energy_array[i] = sample_energy(
@@ -584,7 +588,7 @@ def create_packet_nus(
                 else:
                     energy_array[i] = 511
             else:
-                energy_array[i] = packets.iloc[i]["radiation_energy_kev"]
+                energy_array[i] = packets.iloc[i]["radiation_energy_keV"]
 
         return energy_array
 
@@ -664,13 +668,13 @@ def create_packet_times_uniform_energy(
         array
             Array of decay times
         """
-        decay_times = np.zeros(len(no_of_packets))
+        decay_times = np.zeros(no_of_packets)
         for i, isotope in enumerate(isotopes.to_numpy()):
-            decay_time_min = self.times[decay_time[i]]
+            decay_time_min = self.times[decay_time]
             if decay_time_min == self.times[-1]:
                 decay_time_max = self.effective_times[-1]
             else:
-                decay_time_max = self.times[decay_time[i] + 1]
+                decay_time_max = self.times[decay_time + 1]
             # rejection sampling
             while (decay_times[i] <= decay_time_min) or (
                 decay_times[i] >= decay_time_max
@@ -680,6 +684,30 @@ def create_packet_times_uniform_energy(
                 )
         return decay_times
 
+    def create_packet_decay_times(self, number_of_packets, isotopes):
+        """Samples the decay time from the mean lifetime of the isotopes
+
+        Parameters
+        ----------
+        number_of_packets : int
+            Number of packets
+        times : array
+            Array of time steps
+
+        Returns
+        -------
+        array
+            Array of decay times
+        """
+        decay_times = np.zeros(number_of_packets)
+        for i, isotope in enumerate(isotopes.to_numpy()):
+            if isotope in self.taus:
+                decay_times[i] = -self.taus[isotope] * np.log(
+                    np.random.random()
+                )
+
+        return decay_times
+
     def create_packets(
         self, decays_per_isotope, number_of_packets, *args, **kwargs
     ):
@@ -739,26 +767,37 @@ def create_packets(
         initial_radii = self.create_packet_radii(sampled_packets_df)
 
         # get the time step index of the packets
-        initial_time_indexes = sampled_packets_df.index.get_level_values(0)
+        # Comment: Need to discuss. the indices should be used
+        # We can use the indices of the effective time array used as a timestamp for each packet
 
-        # get the time of the middle of the step for each packet
-        packet_effective_times = np.array(
-            [self.effective_times[i] for i in initial_time_indexes]
+        packet_effective_times = sampled_packets_df.index.get_level_values(0)
+        initial_time_indexes = np.indices(
+            sampled_packets_df.index.get_level_values(0).shape
         )
 
+        # get the time of the middle of the step for each packet
+        # Comment: Array indices need to be integers
+
+        # packet_effective_times = np.array(
+        #    [self.effective_times[i] for i in initial_time_indexes]
+        # )
+
+        # Comment: Taking too much time for the rejection sampling. Can it be improved with numba?
         # packet decay time
-        times = self.create_packet_times_uniform_energy(
-            number_of_packets,
-            sampled_packets_df.index.get_level_values(2),
-            packet_effective_times,
+        decay_times = self.create_packet_decay_times(
+            number_of_packets, sampled_packets_df.index.get_level_values(2)
         )
+        # times_indices = np.indices(times.shape).flatten()
 
         # scale radius by packet decay time. This could be replaced with
         # Geometry object calculations. Note that this also adds a random
         # unit vector multiplication for 3D. May not be needed.
+
+        # Comment: initial radii needs to be an array
+
         locations = (
-            initial_radii
-            * packet_effective_times
+            initial_radii.values
+            * packet_effective_times.values
             * self.create_packet_directions(number_of_packets)
         )
 
@@ -788,35 +827,42 @@ def create_packets(
         # non-relativistic
         packet_energies_rf = np.zeros(number_of_packets)
         nus_rf = np.zeros(number_of_packets)
-        for i in range(number_of_packets):
-            doppler_factor = doppler_factor_3d(
-                directions[:, i],
-                locations[:, i],
-                times[i],
-            )
-            packet_energies_rf[i] = packet_energies_cmf[i] / doppler_factor
-            nus_rf[i] = nus_cmf[i] / doppler_factor
-
-            # deposit positron energy in both output arrays
-            # this is an average across all packets that are created
-            # it could be changed to be only for packets that are from positrons
-            self.energy_plot_positron_rows[i] = np.array(
-                [
-                    i,
-                    isotope_positron_fraction[sampled_packets_df["isotopes"][i]]
-                    * packet_energies_cmf[i],
-                    # this needs to be sqrt(sum of squares) to get radius
-                    np.linalg.norm(locations[i]),
-                    times[i],
-                ]
-            )
 
-            # this is an average across all packets that are created
-            # it could be changed to be only for packets that are from positrons
-            self.energy_df_rows[shells[i], times[i]] += (
-                isotope_positron_fraction[sampled_packets_df["isotopes"][i]]
-                * packet_energies_cmf[i]
-            )
+        doppler_factors = doppler_factor_3D_all_packets(
+            directions, locations, times
+        )
+
+        packet_energies_rf = packet_energies_cmf / doppler_factors
+        nus_rf = nus_cmf / doppler_factors
+        # The for loop makes this program slow. Changed to numpy multiply
+
+        # deposit positron energy in both output arrays
+        # this is an average across all packets that are created
+        # it could be changed to be only for packets that are from positrons
+
+        # self.energy_plot_positron_rows[i] = np.array(
+        #     [
+        #         i,
+        #         isotope_positron_fraction[
+        #             sampled_packets_df.index.get_level_values(2)[i]
+        #         ]
+        #         * packet_energies_cmf[i],
+        #         # this needs to be sqrt(sum of squares) to get radius
+        #         np.linalg.norm(locations[i]),
+        #         times[i],
+        #     ]
+        # )
+
+        # this is an average across all packets that are created
+        # it could be changed to be only for packets that are from positrons
+
+        # comment: we need the indices here
+        # self.energy_df_rows[shells[i], times_indices[i]] += (
+        #     isotope_positron_fraction[
+        #         sampled_packets_df.index.get_level_values(2)[i]
+        #     ]
+        #     * packet_energies_cmf[i]
+        # )
 
         return GXPacketCollection(
             locations,
@@ -846,10 +892,31 @@ def calculate_positron_fraction(self, isotopes):
         """
         positron_fraction = {}
 
+        # Find the positron fraction from the zeroth shell of the dataframe
+        shell_number_0 = self.isotope_decay_df[
+            self.isotope_decay_df.index.get_level_values("shell_number") == 0
+        ]
+
+        positrons_decay_df = shell_number_0[shell_number_0["radiation"] == "bp"]
+        positron_energy_per_isotope = positrons_decay_df.groupby("isotope")[
+            "radiation_energy_keV"
+        ].sum()
+        # Find the total energy released per isotope from the dataframe
+        total_energy_per_isotope = shell_number_0.groupby("isotope")[
+            "radiation_energy_keV"
+        ].sum()
+
         for isotope in isotopes:
-            isotope_energy = self.gamma_ray_lines[isotope][0, :]
-            isotope_intensity = self.gamma_ray_lines[isotope][1, :]
-            positron_fraction[isotope] = self.average_positron_energies[
-                isotope
-            ] / np.sum(isotope_energy * isotope_intensity)
+            positron_fraction[isotope] = (
+                positron_energy_per_isotope[isotope]
+                / total_energy_per_isotope[isotope]
+            )
         return positron_fraction
+
+        # for isotope in isotopes:
+        #    isotope_energy = self.gamma_ray_lines[isotope][0, :]
+        #    isotope_intensity = self.gamma_ray_lines[isotope][1, :]
+        #    positron_fraction[isotope] = self.average_positron_energies[
+        #        isotope
+        #    ] / np.sum(isotope_energy * isotope_intensity)
+        # return positron_fraction
diff --git a/tardis/energy_input/gamma_ray_transport.py b/tardis/energy_input/gamma_ray_transport.py
index 0412195e1e0..eeb843caea1 100644
--- a/tardis/energy_input/gamma_ray_transport.py
+++ b/tardis/energy_input/gamma_ray_transport.py
@@ -210,7 +210,7 @@ def calculate_average_energies(raw_isotope_abundance, gamma_ray_lines):
     average_energies = {}
     average_positron_energies = {}
 
-    for i, isotope in enumerate(all_isotope_names):
+    for _, isotope in enumerate(all_isotope_names):
         energy, intensity = setup_input_energy(
             gamma_ray_lines[gamma_ray_lines.index == isotope.replace("-", "")],
             "g",
diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index 654656300ad..096bf27c1a6 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/tardis/energy_input/main_gamma_ray_loop.py
@@ -19,7 +19,10 @@
     create_isotope_dicts_old,
     create_inventories_dict_old,
 )
-from tardis.energy_input.gamma_ray_packet_source import RadioactivePacketSource
+from tardis.energy_input.gamma_ray_packet_source import (
+    RadioactivePacketSource,
+    GammaRayPacketSource,
+)
 from tardis.energy_input.gamma_ray_transport import (
     calculate_average_energies,
     calculate_average_power_per_mass,
@@ -31,12 +34,13 @@
     iron_group_fraction_per_shell,
 )
 from tardis.energy_input.GXPacket import GXPacket
+from tardis.energy_input.util import make_isotope_string_tardis_like
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
 
-def run_gamma_ray_loop(
+def run_gamma_ray_loop_old(
     model,
     plasma,
     num_decays,
@@ -286,3 +290,193 @@ def run_gamma_ray_loop(
         energy_plot_positrons,
         energy_estimated_deposition,
     )
+
+
+def get_effective_time_array(time_start, time_end, time_space, time_steps):
+    """
+    Function to get the effective time array for the gamma-ray loop.
+
+    Parameters
+    ----------
+    time_start : float
+        start time in days.
+    time_end : float
+        end time in days.
+    time_space : str
+        linear or log.
+    time_steps : int
+        number of time steps.
+
+    Returns
+    -------
+    times : np.ndarray
+        array of times in secs.
+    effective_time_array : np.ndarray
+        effective time array in secs.
+    """
+
+    time_start *= u.d.to(u.s)
+    time_end *= u.d.to(u.s)
+
+    assert time_start < time_end, "time_start must be smaller than time_end!"
+    if time_space == "log":
+        times = np.geomspace(time_start, time_end, time_steps + 1)
+        effective_time_array = np.sqrt(times[:-1] * times[1:])
+    else:
+        times = np.linspace(time_start, time_end, time_steps + 1)
+        effective_time_array = 0.5 * (times[:-1] + times[1:])
+
+    return times, effective_time_array
+
+
+def run_gamma_ray_loop(
+    model,
+    plasma,
+    isotope_decay_df,
+    cumulative_decays_df,
+    num_decays,
+    times,
+    effective_time_array,
+    seed,
+    positronium_fraction,
+    path_to_decay_data,
+    spectrum_bins,
+    time_steps,
+    grey_opacity,
+    photoabsorption_opacity="tardis",
+    pair_creation_opacity="tardis",
+):
+    np.random.seed(seed)
+    inner_velocities = model.v_inner.to("cm/s").value
+    outer_velocities = model.v_outer.to("cm/s").value
+    ejecta_volume = model.volume.to("cm^3").value
+    number_of_shells = model.no_of_shells
+    shell_masses = model.volume * model.density
+    raw_isotope_abundance = model.composition.raw_isotope_abundance.sort_values(
+        by=["atomic_number", "mass_number"], ascending=False
+    )
+
+    dt_array = np.diff(times)
+
+    ejecta_velocity_volume = calculate_ejecta_velocity_volume(model)
+
+    inv_volume_time = (
+        1.0 / ejecta_velocity_volume[:, np.newaxis]
+    ) / effective_time_array**3.0
+
+    for atom_number in plasma.isotope_number_density.index.get_level_values(0):
+        values = plasma.isotope_number_density.loc[atom_number].values
+        if values.shape[0] > 1:
+            plasma.isotope_number_density.loc[atom_number].update = np.sum(
+                values, axis=0
+            )
+        else:
+            plasma.isotope_number_density.loc[atom_number].update = values
+
+    # Electron number density
+    electron_number_density = plasma.number_density.mul(
+        plasma.number_density.index, axis=0
+    ).sum()
+    electron_number = np.array(electron_number_density * ejecta_volume)
+
+    # Evolve electron number and mass density with time
+    electron_number_density_time = (
+        electron_number[:, np.newaxis] * inv_volume_time
+    )
+    mass_density_time = shell_masses[:, np.newaxis] * inv_volume_time
+
+    # gamma_ray_lines = get_nuclear_lines_database(path_to_decay_data)
+    taus, parents = get_taus(raw_isotope_abundance)
+    # Need to get the strings for the isotopes without the dashes
+    taus = make_isotope_string_tardis_like(taus)
+
+    total_energy = cumulative_decays_df[
+        "decay_energy_keV"
+    ].sum()  # total energy in keV
+
+    energy_per_packet = total_energy / num_decays
+    energy_df_rows = np.zeros((number_of_shells, times.shape[0]))
+
+    logger.info(f"Total gamma-ray energy is {total_energy}")
+    average_power_per_mass = 1e41
+    positron_energy_per_isotope = 1000  # keV
+
+    packet_source = GammaRayPacketSource(
+        energy_per_packet,
+        isotope_decay_df,
+        positronium_fraction,
+        inner_velocities,
+        outer_velocities,
+        inv_volume_time,
+        times,
+        energy_df_rows,
+        effective_time_array,
+        taus,
+        parents,
+        positron_energy_per_isotope,
+        average_power_per_mass,
+    )
+    logger.info("Creating packets")
+    packet_collection = packet_source.create_packets(
+        cumulative_decays_df, num_decays
+    )
+    logger.info("Creating packet list")
+    packets = []
+    packets = [
+        GXPacket(
+            packet_collection.location[:, i],
+            packet_collection.direction[:, i],
+            packet_collection.energy_rf[i],
+            packet_collection.energy_cmf[i],
+            packet_collection.nu_rf[i],
+            packet_collection.nu_cmf[i],
+            packet_collection.status[i],
+            packet_collection.shell[i],
+            packet_collection.time_current[i],
+        )
+        for i in range(num_decays)
+    ]
+
+    # return packets
+
+    energy_bins = np.logspace(2, 3.8, spectrum_bins)
+    energy_out = np.zeros((len(energy_bins - 1), time_steps))
+    packets_info_array = np.zeros((int(num_decays), 8))
+    iron_group_fraction = iron_group_fraction_per_shell(model)
+
+    logger.info("Entering the main gamma-ray loop")
+
+    gamma_packet_loop(
+        packets,
+        grey_opacity,
+        photoabsorption_opacity,
+        pair_creation_opacity,
+        electron_number_density_time,
+        mass_density_time,
+        inv_volume_time,
+        iron_group_fraction.to_numpy(),
+        inner_velocities,
+        outer_velocities,
+        times,
+        dt_array,
+        effective_time_array,
+        energy_bins,
+        energy_out,
+        packets_info_array,
+    )
+
+    # packets_df_escaped = pd.DataFrame(
+    #     data=packets_array,
+    #     columns=[
+    #         "packet_index",
+    #         "status",
+    #         "nu_cmf",
+    #         "nu_rf",
+    #         "energy_cmf",
+    #         "lum_rf",
+    #         "energy_rf",
+    #         "shell_number",
+    #     ],
+    # )
+
+    # return packets_df_escaped
diff --git a/tardis/energy_input/util.py b/tardis/energy_input/util.py
index d50f5893d35..f69a78504e0 100644
--- a/tardis/energy_input/util.py
+++ b/tardis/energy_input/util.py
@@ -72,6 +72,28 @@ def doppler_factor_3d(direction_vector, position_vector, time):
     return 1 - (np.dot(direction_vector, velocity_vector) / C_CGS)
 
 
+@njit(**njit_dict_no_parallel)
+def doppler_factor_3D_all_packets(direction_vectors, position_vectors, times):
+    """Doppler shift for photons in 3D
+
+    Parameters
+    ----------
+    direction_vectors : array
+    position_vectors : array
+    times : array
+
+    Returns
+    -------
+    array
+        Doppler factors
+    """
+    velocity_vector = position_vectors / times
+    vel_mul_dir = np.multiply(velocity_vector, direction_vectors)
+    doppler_factors = 1 - (np.sum(vel_mul_dir, axis=0) / C_CGS)
+
+    return doppler_factors
+
+
 @njit(**njit_dict_no_parallel)
 def angle_aberration_gamma(direction_vector, position_vector, time):
     """Angle aberration formula for photons in 3D
@@ -385,3 +407,26 @@ def get_index(value, array):
         i += 1
 
     return i
+
+
+def make_isotope_string_tardis_like(isotope_dict):
+    """Converts isotope string to TARDIS format
+
+    Parameters
+    ----------
+    isotope : str
+        Isotope string
+
+    Returns
+    -------
+    str
+        TARDIS-like isotope string
+    """
+
+    new_isotope_dict = {}
+
+    for key in isotope_dict.keys():
+        new_key = key.replace("-", "")
+        new_isotope_dict[new_key] = isotope_dict[key]
+
+    return new_isotope_dict
diff --git a/tardis/model/matter/decay.py b/tardis/model/matter/decay.py
index 5d5e00279dc..43cff632cd0 100644
--- a/tardis/model/matter/decay.py
+++ b/tardis/model/matter/decay.py
@@ -95,9 +95,7 @@ def decay(self, t):
             Decayed abundances
         """
         inventories = self.to_inventories()
-        t_second = (
-            u.Quantity(t, u.day).to(u.s).value - self.time_0.to(u.s).value
-        )
+        t_second = u.Quantity(t, u.s).to(u.s).value - self.time_0.to(u.s).value
         logger.info(f"Decaying abundances for {t_second} seconds")
         if t_second < 0:
             logger.warning(

From 28f45545974b005068f9f0700e68c7f954bf775a Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Mon, 6 May 2024 15:43:57 -0400
Subject: [PATCH 03/18] Added a rejection sampling into decay times

---
 tardis/energy_input/gamma_packet_loop.py      | 64 ++++++++++++++++-
 tardis/energy_input/gamma_ray_grid.py         |  5 +-
 .../energy_input/gamma_ray_packet_source.py   | 68 ++++++++++++++-----
 tardis/energy_input/main_gamma_ray_loop.py    | 59 +++++++++++-----
 4 files changed, 160 insertions(+), 36 deletions(-)

diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py
index 17f93775d44..7c0fc7238a4 100644
--- a/tardis/energy_input/gamma_packet_loop.py
+++ b/tardis/energy_input/gamma_packet_loop.py
@@ -15,6 +15,7 @@
 from tardis.energy_input.gamma_ray_grid import (
     distance_trace,
     move_packet,
+    calculate_distance_radial,
 )
 from tardis.energy_input.util import (
     doppler_factor_3d,
@@ -112,7 +113,7 @@ def gamma_packet_loop(
 
     for i in range(packet_count):
         packet = packets[i]
-        time_index = get_index(packet.time_current, times)
+        time_index = get_index(packet.time_current, effective_time_array)
         # print("Time index: ", time_index)
         # print("Time current: ", packet.time_current)
 
@@ -376,3 +377,64 @@ def process_packet_path(packet):
         ejecta_energy_gained = packet.energy_cmf
 
     return packet, ejecta_energy_gained
+
+
+@njit(**njit_dict_no_parallel)
+def gamma_packet_loop_understanding(
+    packets,
+    times,
+    inner_velocity,
+    outer_velocity,
+    effective_time_array,
+):
+
+    packet_count = len(packets)
+    print("Packet count: ", packet_count)
+    print("Effective time array: ", effective_time_array)
+    print("length of effective time array", len(effective_time_array))
+
+    for i in range(packet_count):
+        packet = packets[i]
+        time_index = get_index(packet.time_current, effective_time_array)
+        print("Time index: ", time_index)
+        print("Time current: ", packet.time_current)
+        print("Time current in days: ", packet.time_current / 86400)
+        print("Length of time array:", len(times))
+
+        # we need to have the times within the simulation time steps
+
+        if time_index < 0:
+            print(packet.time_current, time_index)
+            raise ValueError("Packet time index less than 0!")
+        elif time_index >= len(times):
+            print(packet.time_current, time_index)
+            raise ValueError("Packet time index greater than time steps!")
+
+        # while packet.status == GXPacketStatus.IN_PROCESS:
+
+        print(
+            "Inner shell velocity: ",
+            inner_velocity[packet.shell],
+        )
+        print(
+            "Outer shell velocity: ",
+            outer_velocity[packet.shell],
+        )
+
+        print(
+            "Effective time array: ",
+            effective_time_array[time_index],
+        )
+        print(
+            "Effective time array in days:",
+            effective_time_array[time_index] / 86400,
+        )
+
+        distance_boundary, shell_change = calculate_distance_radial(
+            packet,
+            inner_velocity[packet.shell] * packet.time_current,
+            outer_velocity[packet.shell] * packet.time_current,
+        )
+
+        print("Distance boundary: ", distance_boundary)
+        print("Shell change: ", shell_change)
diff --git a/tardis/energy_input/gamma_ray_grid.py b/tardis/energy_input/gamma_ray_grid.py
index 56f7a939a76..747ae63e8cc 100644
--- a/tardis/energy_input/gamma_ray_grid.py
+++ b/tardis/energy_input/gamma_ray_grid.py
@@ -53,6 +53,7 @@ def calculate_distance_radial(photon, r_inner, r_outer):
 
     # the correct distance is the shortest positive distance
     distance_list = [i for i in distances if i > 0]
+    print(distance_list)
 
     if not distance_list:
         print(photon.get_location_r() - r_inner)
@@ -103,8 +104,8 @@ def distance_trace(
     """
     distance_boundary, shell_change = calculate_distance_radial(
         photon,
-        inner_velocity[photon.shell] * current_time,
-        outer_velocity[photon.shell] * current_time,
+        inner_velocity[photon.shell] * photon.time_current,
+        outer_velocity[photon.shell] * photon.time_current,
     )
 
     distance_interaction = photon.tau / total_opacity
diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index 1bcc008a4e1..6432000b214 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -708,6 +708,45 @@ def create_packet_decay_times(self, number_of_packets, isotopes):
 
         return decay_times
 
+    def sample_decay_times(self, isotopes, taus, times, number_of_packets):
+        """
+        Sample decay times for a given number of packets
+
+        Parameters
+        ----------
+        isotopes : array
+            Array of isotope names as strings
+        taus : dict
+            Dictionary of isotope mean lifetimes in seconds
+        times : array
+            Array of time steps in seconds
+        number_of_packets : int
+            Number of packets
+
+        Returns
+        -------
+        decay_times : array
+            Array of decay times for each packet
+
+        """
+
+        decay_times = np.zeros(number_of_packets)
+
+        for i, isotope in enumerate(isotopes):
+            # decay time of the packet cannot be less than the minimum time in the simulation
+            decay_time_min = times[0]
+            # decay time of the packet cannot be greater than the maximum time in the simulation
+            decay_time_max = times[-1]
+
+            # rejection sampling on decay times
+            while (decay_times[i] <= decay_time_min) or (
+                decay_times[i] >= decay_time_max
+            ):
+                # since we know which isotope the packet is associated with, we can use the tau value for that isotope
+                decay_times[i] = -taus[isotope] * np.log(np.random.random())
+
+        return decay_times
+
     def create_packets(
         self, decays_per_isotope, number_of_packets, *args, **kwargs
     ):
@@ -750,7 +789,7 @@ def create_packets(
             random_state=np.random.RandomState(self.base_seed),
         )
         # get unique isotopes that have produced packets
-        isotopes = pd.unique(sampled_packets_df.index.get_level_values(2))
+        isotopes = sampled_packets_df.index.get_level_values(2)
 
         # compute the positron fraction for unique isotopes
         isotope_positron_fraction = self.calculate_positron_fraction(isotopes)
@@ -770,23 +809,13 @@ def create_packets(
         # Comment: Need to discuss. the indices should be used
         # We can use the indices of the effective time array used as a timestamp for each packet
 
-        packet_effective_times = sampled_packets_df.index.get_level_values(0)
-        initial_time_indexes = np.indices(
-            sampled_packets_df.index.get_level_values(0).shape
-        )
+        # packet_effective_times = sampled_packets_df.index.get_level_values(0)
+        # initial_time_indexes = np.indices(
+        #    sampled_packets_df.index.get_level_values(0).shape
+        # )
 
         # get the time of the middle of the step for each packet
         # Comment: Array indices need to be integers
-
-        # packet_effective_times = np.array(
-        #    [self.effective_times[i] for i in initial_time_indexes]
-        # )
-
-        # Comment: Taking too much time for the rejection sampling. Can it be improved with numba?
-        # packet decay time
-        decay_times = self.create_packet_decay_times(
-            number_of_packets, sampled_packets_df.index.get_level_values(2)
-        )
         # times_indices = np.indices(times.shape).flatten()
 
         # scale radius by packet decay time. This could be replaced with
@@ -795,9 +824,14 @@ def create_packets(
 
         # Comment: initial radii needs to be an array
 
+        packet_effective_times = sampled_packets_df.index.get_level_values(0)
+        decay_times = self.sample_decay_times(
+            isotopes, self.taus, self.effective_times, number_of_packets
+        )
+
         locations = (
             initial_radii.values
-            * packet_effective_times.values
+            * decay_times
             * self.create_packet_directions(number_of_packets)
         )
 
@@ -873,7 +907,7 @@ def create_packets(
             nus_cmf,
             statuses,
             shells,
-            times,
+            decay_times,
         )
 
     def calculate_positron_fraction(self, isotopes):
diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index 096bf27c1a6..b8412826f0c 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/tardis/energy_input/main_gamma_ray_loop.py
@@ -7,7 +7,10 @@
 from tardis.energy_input.energy_source import (
     get_nuclear_lines_database,
 )
-from tardis.energy_input.gamma_packet_loop import gamma_packet_loop
+from tardis.energy_input.gamma_packet_loop import (
+    gamma_packet_loop,
+    gamma_packet_loop_understanding,
+)
 from tardis.energy_input.gamma_ray_channel import (
     calculate_total_decays,
     create_inventories_dict,
@@ -446,7 +449,31 @@ def run_gamma_ray_loop(
 
     logger.info("Entering the main gamma-ray loop")
 
-    gamma_packet_loop(
+    # packets_df_escaped = pd.DataFrame(
+    #     data=packets_array,
+    #     columns=[
+    #         "packet_index",
+    #         "status",
+    #         "nu_cmf",
+    #         "nu_rf",
+    #         "energy_cmf",
+    #         "lum_rf",
+    #         "energy_rf",
+    #         "shell_number",
+    #     ],
+    # )
+
+    # return packets_df_escaped
+
+    # gamma_packet_loop_understanding(
+    #     packets, times, inner_velocities, outer_velocities, effective_time_array
+    # )
+
+    (
+        energy_out,
+        bin_width,
+        packets_array,
+    ) = gamma_packet_loop(
         packets,
         grey_opacity,
         photoabsorption_opacity,
@@ -465,18 +492,18 @@ def run_gamma_ray_loop(
         packets_info_array,
     )
 
-    # packets_df_escaped = pd.DataFrame(
-    #     data=packets_array,
-    #     columns=[
-    #         "packet_index",
-    #         "status",
-    #         "nu_cmf",
-    #         "nu_rf",
-    #         "energy_cmf",
-    #         "lum_rf",
-    #         "energy_rf",
-    #         "shell_number",
-    #     ],
-    # )
+    packets_df_escaped = pd.DataFrame(
+        data=packets_array,
+        columns=[
+            "packet_index",
+            "status",
+            "nu_cmf",
+            "nu_rf",
+            "energy_cmf",
+            "lum_rf",
+            "energy_rf",
+            "shell_number",
+        ],
+    )
 
-    # return packets_df_escaped
+    return packets_df_escaped

From 64ab1f007eff7fff9863579bd4ec4cfcfeee4abb Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Tue, 14 May 2024 17:08:20 -0400
Subject: [PATCH 04/18] Changed gamma ray grid

---
 tardis/energy_input/gamma_ray_grid.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tardis/energy_input/gamma_ray_grid.py b/tardis/energy_input/gamma_ray_grid.py
index 747ae63e8cc..f7f66134725 100644
--- a/tardis/energy_input/gamma_ray_grid.py
+++ b/tardis/energy_input/gamma_ray_grid.py
@@ -53,7 +53,6 @@ def calculate_distance_radial(photon, r_inner, r_outer):
 
     # the correct distance is the shortest positive distance
     distance_list = [i for i in distances if i > 0]
-    print(distance_list)
 
     if not distance_list:
         print(photon.get_location_r() - r_inner)

From 6fb737c8e16c6f543c85dd3a067da8c3b7ae06d0 Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Wed, 12 Jun 2024 13:54:09 -0400
Subject: [PATCH 05/18] modified the input to gammaraypacketsource

---
 tardis/energy_input/gamma_packet_loop.py       |  8 ++++----
 tardis/energy_input/gamma_ray_grid.py          | 13 +++++++++++--
 tardis/energy_input/gamma_ray_packet_source.py | 17 +++++++++++++++--
 tardis/energy_input/main_gamma_ray_loop.py     |  6 +++++-
 4 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py
index 7c0fc7238a4..39046ef931e 100644
--- a/tardis/energy_input/gamma_packet_loop.py
+++ b/tardis/energy_input/gamma_packet_loop.py
@@ -113,7 +113,7 @@ def gamma_packet_loop(
 
     for i in range(packet_count):
         packet = packets[i]
-        time_index = get_index(packet.time_current, effective_time_array)
+        time_index = get_index(packet.time_current, times)
         # print("Time index: ", time_index)
         # print("Time current: ", packet.time_current)
 
@@ -395,7 +395,7 @@ def gamma_packet_loop_understanding(
 
     for i in range(packet_count):
         packet = packets[i]
-        time_index = get_index(packet.time_current, effective_time_array)
+        time_index = get_index(packet.time_current, times)
         print("Time index: ", time_index)
         print("Time current: ", packet.time_current)
         print("Time current in days: ", packet.time_current / 86400)
@@ -432,8 +432,8 @@ def gamma_packet_loop_understanding(
 
         distance_boundary, shell_change = calculate_distance_radial(
             packet,
-            inner_velocity[packet.shell] * packet.time_current,
-            outer_velocity[packet.shell] * packet.time_current,
+            inner_velocity[packet.shell] * effective_time_array[time_index],
+            outer_velocity[packet.shell] * effective_time_array[time_index],
         )
 
         print("Distance boundary: ", distance_boundary)
diff --git a/tardis/energy_input/gamma_ray_grid.py b/tardis/energy_input/gamma_ray_grid.py
index f7f66134725..18829b61d22 100644
--- a/tardis/energy_input/gamma_ray_grid.py
+++ b/tardis/energy_input/gamma_ray_grid.py
@@ -51,6 +51,15 @@ def calculate_distance_radial(photon, r_inner, r_outer):
 
     distances = np.array([inner_1, inner_2, outer_1, outer_2])
 
+    # print(
+    #     "Distance to inner shell boundaries: ",
+    #     photon.get_location_r() - r_inner,
+    # )
+    # print(
+    #     "Distance to outer shell boundaries: ",
+    #     photon.get_location_r() - r_outer,
+    # )
+
     # the correct distance is the shortest positive distance
     distance_list = [i for i in distances if i > 0]
 
@@ -103,8 +112,8 @@ def distance_trace(
     """
     distance_boundary, shell_change = calculate_distance_radial(
         photon,
-        inner_velocity[photon.shell] * photon.time_current,
-        outer_velocity[photon.shell] * photon.time_current,
+        inner_velocity[photon.shell] * current_time,
+        outer_velocity[photon.shell] * current_time,
     )
 
     distance_interaction = photon.tau / total_opacity
diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index 6432000b214..b99cb188ce0 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -824,14 +824,27 @@ def create_packets(
 
         # Comment: initial radii needs to be an array
 
-        packet_effective_times = sampled_packets_df.index.get_level_values(0)
+        # packet_effective_times = sampled_packets_df.index.get_level_values(0)
+
         decay_times = self.sample_decay_times(
             isotopes, self.taus, self.effective_times, number_of_packets
         )
 
+        # decay_time_indices = np.digitize(
+        #    decay_times, sorted(packet_effective_times.values), right=False
+        # )
+        decay_time_indices = []
+        for i in range(number_of_packets):
+            decay_time_indices.append(
+                get_index(decay_times[i], self.effective_times)
+            )
+        # decay_times = self.create_packet_decay_times(
+        #     number_of_packets, isotopes
+        # )
+
         locations = (
             initial_radii.values
-            * decay_times
+            * self.effective_times[decay_time_indices]
             * self.create_packet_directions(number_of_packets)
         )
 
diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index b8412826f0c..77f3fb776ab 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/tardis/energy_input/main_gamma_ray_loop.py
@@ -506,4 +506,8 @@ def run_gamma_ray_loop(
         ],
     )
 
-    return packets_df_escaped
+    escape_energy = pd.DataFrame(
+        data=energy_out, columns=times[:-1], index=energy_bins
+    )
+
+    return escape_energy, packets_df_escaped

From cc9fe352ea9dda438af39a1ee2df043ed66b2edf Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Thu, 13 Jun 2024 11:37:14 -0400
Subject: [PATCH 06/18] Added dataframe for storing packet info

---
 tardis/energy_input/gamma_packet_loop.py      | 14 +++++++----
 .../energy_input/gamma_ray_packet_source.py   |  8 +++++--
 tardis/energy_input/main_gamma_ray_loop.py    | 23 +++++++++++++++++--
 3 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py
index 39046ef931e..e4581f76eaf 100644
--- a/tardis/energy_input/gamma_packet_loop.py
+++ b/tardis/energy_input/gamma_packet_loop.py
@@ -21,6 +21,7 @@
     doppler_factor_3d,
     C_CGS,
     H_CGS_KEV,
+    KEV2ERG,
     get_index,
 )
 from tardis.energy_input.GXPacket import GXPacketStatus
@@ -209,7 +210,7 @@ def gamma_packet_loop(
                 outer_velocities,
                 total_opacity,
                 effective_time_array[time_index],
-                times[time_index + 1],
+                effective_time_array[time_index + 1],
             )
 
             distance = min(
@@ -285,14 +286,17 @@ def gamma_packet_loop(
 
                 if packet.shell > len(mass_density_time[:, 0]) - 1:
                     rest_energy = packet.nu_rf * H_CGS_KEV
-                    lum_rf = (packet.energy_rf * 1.6022e-9) / dt
                     bin_index = get_index(rest_energy, energy_bins)
                     bin_width = (
                         energy_bins[bin_index + 1] - energy_bins[bin_index]
                     )
-                    energy_out[bin_index, time_index] += rest_energy / (
-                        bin_width * dt
+                    freq_bin_width = bin_width / H_CGS_KEV
+                    energy_ergs = packet.energy_rf * KEV2ERG
+                    luminosity = energy_ergs / dt
+                    energy_out[bin_index, time_index] += (
+                        energy_ergs / dt / freq_bin_width
                     )
+
                     packet.status = GXPacketStatus.ESCAPED
                     escaped_packets += 1
                     if scattered:
@@ -309,7 +313,7 @@ def gamma_packet_loop(
                     packet.nu_cmf,
                     packet.nu_rf,
                     packet.energy_cmf,
-                    lum_rf,
+                    luminosity,
                     packet.energy_rf,
                     packet.shell,
                 ]
diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index b99cb188ce0..48d8514f726 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -782,7 +782,11 @@ def create_packets(
 
         # sample packets from dataframe, returning a dataframe where each row is
         # a sampled packet
-        sampled_packets_df = decays_per_isotope.sample(
+
+        sampled_packets_df_gamma = decays_per_isotope[
+            decays_per_isotope["radiation"] == "g"
+        ]
+        sampled_packets_df = sampled_packets_df_gamma.sample(
             n=number_of_packets,
             weights="decay_energy_erg",
             replace=True,
@@ -876,7 +880,7 @@ def create_packets(
         nus_rf = np.zeros(number_of_packets)
 
         doppler_factors = doppler_factor_3D_all_packets(
-            directions, locations, times
+            directions, locations, decay_times
         )
 
         packet_energies_rf = packet_energies_cmf / doppler_factors
diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index 77f3fb776ab..167d94e47c2 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/tardis/energy_input/main_gamma_ray_loop.py
@@ -393,14 +393,21 @@ def run_gamma_ray_loop(
     # Need to get the strings for the isotopes without the dashes
     taus = make_isotope_string_tardis_like(taus)
 
-    total_energy = cumulative_decays_df[
+    cumulative_decays_df_gamma = cumulative_decays_df[
+        cumulative_decays_df["radiation"] == "g"
+    ]
+
+    total_energy = cumulative_decays_df_gamma[
         "decay_energy_keV"
     ].sum()  # total energy in keV
 
+    # total_energy_erg = total_energy * u.keV.to("erg")
+
     energy_per_packet = total_energy / num_decays
     energy_df_rows = np.zeros((number_of_shells, times.shape[0]))
 
     logger.info(f"Total gamma-ray energy is {total_energy}")
+    logger.info(f"Energy per packet is {energy_per_packet}")
     average_power_per_mass = 1e41
     positron_energy_per_isotope = 1000  # keV
 
@@ -449,6 +456,18 @@ def run_gamma_ray_loop(
 
     logger.info("Entering the main gamma-ray loop")
 
+    total_cmf_energy = 0
+    total_rf_energy = 0
+
+    for p in packets:
+        total_cmf_energy += p.energy_cmf
+        total_rf_energy += p.energy_rf
+
+    print("Total CMF energy")
+    print(total_cmf_energy)
+    print("Total RF energy")
+    print(total_rf_energy)
+
     # packets_df_escaped = pd.DataFrame(
     #     data=packets_array,
     #     columns=[
@@ -500,7 +519,7 @@ def run_gamma_ray_loop(
             "nu_cmf",
             "nu_rf",
             "energy_cmf",
-            "lum_rf",
+            "luminosity",
             "energy_rf",
             "shell_number",
         ],

From b523946f6655bc4943f621b3f91cf94557231f2d Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Thu, 13 Jun 2024 11:51:58 -0400
Subject: [PATCH 07/18] Checked with black

---
 .../energy_input/gamma_ray_packet_source.py   | 24 +++++++++----------
 tardis/energy_input/main_gamma_ray_loop.py    |  6 +----
 2 files changed, 13 insertions(+), 17 deletions(-)

diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index 2fcf2754026..7ff53183b46 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -438,18 +438,18 @@ def create_packets(self, decays_per_isotope, *args, **kwargs):
                 packet_energies_cmf[
                     packet_index - isotope_packet_count : packet_index
                 ] = initial_packet_energies_cmf
-                nus_rf[packet_index - isotope_packet_count : packet_index] = (
-                    initial_nus_rf
-                )
-                nus_cmf[packet_index - isotope_packet_count : packet_index] = (
-                    initial_nus_cmf
-                )
-                shells[packet_index - isotope_packet_count : packet_index] = (
-                    initial_shells
-                )
-                times[packet_index - isotope_packet_count : packet_index] = (
-                    initial_times
-                )
+                nus_rf[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_nus_rf
+                nus_cmf[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_nus_cmf
+                shells[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_shells
+                times[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_times
 
         return GXPacketCollection(
             locations,
diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index 167d94e47c2..5870fcf6a07 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/tardis/energy_input/main_gamma_ray_loop.py
@@ -488,11 +488,7 @@ def run_gamma_ray_loop(
     #     packets, times, inner_velocities, outer_velocities, effective_time_array
     # )
 
-    (
-        energy_out,
-        bin_width,
-        packets_array,
-    ) = gamma_packet_loop(
+    (energy_out, bin_width, packets_array,) = gamma_packet_loop(
         packets,
         grey_opacity,
         photoabsorption_opacity,

From 14ab57eb958e1cb4fc2b46cd6e809045db8b2b86 Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Tue, 18 Jun 2024 16:27:07 -0400
Subject: [PATCH 08/18] added tests for total decay energy

---
 tardis/energy_input/gamma_ray_channel.py      | 12 +--
 .../energy_input/gamma_ray_packet_source.py   |  4 +-
 tardis/energy_input/main_gamma_ray_loop.py    | 41 +++++-----
 .../tests/test_gamma_ray_channel.py           | 79 +++++++++++++++++++
 tardis/model/matter/decay.py                  |  4 +-
 tardis/model/parse_input.py                   |  2 +-
 6 files changed, 114 insertions(+), 28 deletions(-)

diff --git a/tardis/energy_input/gamma_ray_channel.py b/tardis/energy_input/gamma_ray_channel.py
index 42df619d4ed..7321923d90e 100644
--- a/tardis/energy_input/gamma_ray_channel.py
+++ b/tardis/energy_input/gamma_ray_channel.py
@@ -76,7 +76,7 @@ def calculate_total_decays(inventories, time_delta):
         dictionary of inventories for each shell
 
     time_end : float
-        End time of simulation in days.
+        End time of simulation step in days.
 
 
     Returns
@@ -84,10 +84,10 @@ def calculate_total_decays(inventories, time_delta):
     cumulative_decay_df : pd.DataFrame
         total decays for x g of isotope for time 't'
     """
-    time_delta = u.Quantity(time_delta, u.s)
+    time_delta = u.Quantity(time_delta, u.d)
     total_decays = {}
     for shell, inventory in inventories.items():
-        total_decays[shell] = inventory.cumulative_decays(time_delta.value)
+        total_decays[shell] = inventory.cumulative_decays(time_delta.value, "d")
 
     flattened_dict = {}
 
@@ -205,7 +205,7 @@ def evolve_mass_fraction(raw_isotope_mass_fraction, time_array):
     return time_evolved_isotope_mass_fraction
 
 
-def time_evolve_cumulative_decays(
+def time_evolve_mass_fractions(
     raw_isotope_mass_fraction, shell_masses, gamma_ray_lines, time_array
 ):
     """
@@ -234,7 +234,9 @@ def time_evolve_cumulative_decays(
     cumulative_decay_df_list = []
     initial_isotope_mass_fraction = raw_isotope_mass_fraction
 
-    for time in time_array:
+    decay_times = np.diff(time_array)
+
+    for time in decay_times:
         decayed_isotope_mass_fraction = IsotopicMassFraction(
             initial_isotope_mass_fraction
         ).decay(time)
diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index 7ff53183b46..0fa4db9dc31 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -792,13 +792,13 @@ def create_packets(
             random_state=np.random.RandomState(self.base_seed),
         )
         # get unique isotopes that have produced packets
-        isotopes = sampled_packets_df.index.get_level_values(2)
+        isotopes = sampled_packets_df.index.get_level_values(1)
 
         # compute the positron fraction for unique isotopes
         isotope_positron_fraction = self.calculate_positron_fraction(isotopes)
 
         # get the packet shell index
-        shells = sampled_packets_df.index.get_level_values(1)
+        shells = sampled_packets_df.index.get_level_values(0)
 
         # get the inner and outer velocity boundaries for each packet to compute
         # the initial radii
diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index 5870fcf6a07..c57ca720c29 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/tardis/energy_input/main_gamma_ray_loop.py
@@ -318,9 +318,6 @@ def get_effective_time_array(time_start, time_end, time_space, time_steps):
         effective time array in secs.
     """
 
-    time_start *= u.d.to(u.s)
-    time_end *= u.d.to(u.s)
-
     assert time_start < time_end, "time_start must be smaller than time_end!"
     if time_space == "log":
         times = np.geomspace(time_start, time_end, time_steps + 1)
@@ -334,7 +331,6 @@ def get_effective_time_array(time_start, time_end, time_space, time_steps):
 
 def run_gamma_ray_loop(
     model,
-    plasma,
     isotope_decay_df,
     cumulative_decays_df,
     num_decays,
@@ -367,18 +363,25 @@ def run_gamma_ray_loop(
         1.0 / ejecta_velocity_volume[:, np.newaxis]
     ) / effective_time_array**3.0
 
-    for atom_number in plasma.isotope_number_density.index.get_level_values(0):
-        values = plasma.isotope_number_density.loc[atom_number].values
+    for (
+        atom_number
+    ) in model.composition.isotopic_number_density.index.get_level_values(0):
+        values = model.composition.isotopic_number_density.loc[
+            atom_number
+        ].values
         if values.shape[0] > 1:
-            plasma.isotope_number_density.loc[atom_number].update = np.sum(
-                values, axis=0
-            )
+            model.composition.isotopic_number_density.loc[
+                atom_number
+            ].update = np.sum(values, axis=0)
         else:
-            plasma.isotope_number_density.loc[atom_number].update = values
+            model.composition.isotopic_number_density.loc[
+                atom_number
+            ].update = values
 
     # Electron number density
-    electron_number_density = plasma.number_density.mul(
-        plasma.number_density.index, axis=0
+    electron_number_density = model.composition.isotopic_number_density.mul(
+        model.composition.isotopic_number_density.index.get_level_values(0),
+        axis=0,
     ).sum()
     electron_number = np.array(electron_number_density * ejecta_volume)
 
@@ -393,20 +396,20 @@ def run_gamma_ray_loop(
     # Need to get the strings for the isotopes without the dashes
     taus = make_isotope_string_tardis_like(taus)
 
-    cumulative_decays_df_gamma = cumulative_decays_df[
-        cumulative_decays_df["radiation"] == "g"
+    isotope_decay_df_gamma = isotope_decay_df[
+        isotope_decay_df["radiation"] == "g"
     ]
 
-    total_energy = cumulative_decays_df_gamma[
+    total_energy_gamma = isotope_decay_df_gamma[
         "decay_energy_keV"
     ].sum()  # total energy in keV
 
-    # total_energy_erg = total_energy * u.keV.to("erg")
+    total_energy_gamma_erg = total_energy_gamma * u.keV.to("erg")
 
-    energy_per_packet = total_energy / num_decays
+    energy_per_packet = total_energy_gamma_erg / num_decays
     energy_df_rows = np.zeros((number_of_shells, times.shape[0]))
 
-    logger.info(f"Total gamma-ray energy is {total_energy}")
+    logger.info(f"Total energy in gamma-rays is {total_energy_gamma_erg}")
     logger.info(f"Energy per packet is {energy_per_packet}")
     average_power_per_mass = 1e41
     positron_energy_per_isotope = 1000  # keV
@@ -428,7 +431,7 @@ def run_gamma_ray_loop(
     )
     logger.info("Creating packets")
     packet_collection = packet_source.create_packets(
-        cumulative_decays_df, num_decays
+        isotope_decay_df, num_decays
     )
     logger.info("Creating packet list")
     packets = []
diff --git a/tardis/energy_input/tests/test_gamma_ray_channel.py b/tardis/energy_input/tests/test_gamma_ray_channel.py
index 5842ba67783..ea8e16d034b 100644
--- a/tardis/energy_input/tests/test_gamma_ray_channel.py
+++ b/tardis/energy_input/tests/test_gamma_ray_channel.py
@@ -18,6 +18,8 @@
     calculate_total_decays,
     create_isotope_decay_df,
 )
+from tardis.energy_input.gamma_ray_transport import get_taus
+from tardis.energy_input.util import KEV2ERG
 
 
 @pytest.fixture(scope="module")
@@ -234,3 +236,80 @@ def test_activity(gamma_ray_test_composition, nuclide_name):
     expected = number_of_atoms * (1 - np.exp(-decay_constant * time_delta))
 
     npt.assert_allclose(actual, expected)
+
+
+def test_total_energy_production(gamma_ray_test_composition, atomic_dataset):
+    """
+    Function to test the total energy production with equation 18 of Nadyozhin 1994. This is only for Ni56 now.
+    Parameters
+    ----------
+    gamma_ray_test_composition: Function holding the composition.
+    """
+    time_start = 0.0 * u.d
+    time_end = np.inf * u.d
+    time_delta = (time_end - time_start).value
+
+    gamma_ray_lines = atomic_dataset.decay_radiation_data
+    raw_isotopic_mass_fraction, cell_masses = gamma_ray_test_composition
+    isotope_dict = create_isotope_dicts(raw_isotopic_mass_fraction, cell_masses)
+    inventories_dict = create_inventories_dict(isotope_dict)
+    total_decays = calculate_total_decays(inventories_dict, time_delta)
+    isotope_decay_df = create_isotope_decay_df(total_decays, gamma_ray_lines)
+    taus, parents = get_taus(raw_isotopic_mass_fraction)
+
+    ni56_nuclide = rd.Nuclide("Ni56")
+    atomic_mass_unit = const.u.cgs.value
+    tau_ni56 = taus["Ni-56"]
+    tau_co56 = taus["Co-56"]
+
+    ni56_mass_fraction = raw_isotopic_mass_fraction.loc[(28, 56)]
+
+    ni_56_mass = sum(ni56_mass_fraction * cell_masses)
+    ni_56_mass_solar = ni_56_mass / const.M_sun.cgs.value
+
+    shell_number_0 = isotope_decay_df[
+        isotope_decay_df.index.get_level_values("shell_number") == 0
+    ]
+
+    ni56 = shell_number_0[shell_number_0.index.get_level_values(1) == "Ni56"]
+    ni_56_energy = ni56["energy_per_channel_keV"].sum()
+
+    co_56 = shell_number_0[shell_number_0.index.get_level_values(1) == "Co56"]
+    co_56_energy = co_56["energy_per_channel_keV"].sum()
+
+    first_term = const.M_sun.cgs.value / (
+        (tau_co56 - tau_ni56) * ni56_nuclide.atomic_mass * atomic_mass_unit
+    )
+    ni_term = (
+        (ni_56_energy * (tau_co56 / tau_ni56 - 1) - co_56_energy)
+        * first_term
+        * KEV2ERG
+    )
+    co_term = co_56_energy * first_term * KEV2ERG
+
+    expected = (
+        ni_term
+        * tau_ni56
+        * (
+            np.exp(-time_start.value / tau_ni56)
+            - np.exp(-time_end.value / tau_ni56)
+        )
+        + co_term
+        * tau_co56
+        * (
+            np.exp(-time_start.value / tau_co56)
+            - np.exp(-time_end.value / tau_co56)
+        )
+    ) * ni_56_mass_solar
+
+    ni56_df = isotope_decay_df[
+        isotope_decay_df.index.get_level_values(1) == "Ni56"
+    ]
+    ni56_energy = ni56_df["decay_energy_erg"].sum()
+    co_56_df = isotope_decay_df[
+        isotope_decay_df.index.get_level_values(1) == "Co56"
+    ]
+    co56_energy = co_56_df["decay_energy_erg"].sum()
+    actual = ni56_energy + co56_energy
+
+    npt.assert_allclose(actual, expected)
diff --git a/tardis/model/matter/decay.py b/tardis/model/matter/decay.py
index 43cff632cd0..5d5e00279dc 100644
--- a/tardis/model/matter/decay.py
+++ b/tardis/model/matter/decay.py
@@ -95,7 +95,9 @@ def decay(self, t):
             Decayed abundances
         """
         inventories = self.to_inventories()
-        t_second = u.Quantity(t, u.s).to(u.s).value - self.time_0.to(u.s).value
+        t_second = (
+            u.Quantity(t, u.day).to(u.s).value - self.time_0.to(u.s).value
+        )
         logger.info(f"Decaying abundances for {t_second} seconds")
         if t_second < 0:
             logger.warning(
diff --git a/tardis/model/parse_input.py b/tardis/model/parse_input.py
index d1b5ba3766f..340b1af5690 100644
--- a/tardis/model/parse_input.py
+++ b/tardis/model/parse_input.py
@@ -490,7 +490,7 @@ def parse_density_csvy(csvy_model_config, csvy_model_data, time_explosion):
         )
         density_0 = csvy_model_data["density"].values * density_unit
         # Removing as thee new architecture removes the 0th shell already
-        # density_0 = density_0.to("g/cm^3")[1:]
+        density_0 = density_0.to("g/cm^3")[1:]
         # density_0 = density_0.insert(0, 0)
         density = calculate_density_after_time(
             density_0, time_0, time_explosion

From 0a3a1789c96f0b6735dd5e5ed451e598258833e6 Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Thu, 20 Jun 2024 10:27:47 -0400
Subject: [PATCH 09/18] removing accessing any variables in plasma

---
 .../energy_input/gamma_ray_packet_source.py   | 38 +++++++++----------
 tardis/energy_input/main_gamma_ray_loop.py    | 25 ++++++------
 2 files changed, 32 insertions(+), 31 deletions(-)

diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index 0fa4db9dc31..8fb1ecc8927 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -438,18 +438,18 @@ def create_packets(self, decays_per_isotope, *args, **kwargs):
                 packet_energies_cmf[
                     packet_index - isotope_packet_count : packet_index
                 ] = initial_packet_energies_cmf
-                nus_rf[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_nus_rf
-                nus_cmf[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_nus_cmf
-                shells[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_shells
-                times[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_times
+                nus_rf[packet_index - isotope_packet_count : packet_index] = (
+                    initial_nus_rf
+                )
+                nus_cmf[packet_index - isotope_packet_count : packet_index] = (
+                    initial_nus_cmf
+                )
+                shells[packet_index - isotope_packet_count : packet_index] = (
+                    initial_shells
+                )
+                times[packet_index - isotope_packet_count : packet_index] = (
+                    initial_times
+                )
 
         return GXPacketCollection(
             locations,
@@ -792,13 +792,13 @@ def create_packets(
             random_state=np.random.RandomState(self.base_seed),
         )
         # get unique isotopes that have produced packets
-        isotopes = sampled_packets_df.index.get_level_values(1)
+        isotopes = sampled_packets_df.index.get_level_values(2)
 
         # compute the positron fraction for unique isotopes
         isotope_positron_fraction = self.calculate_positron_fraction(isotopes)
 
         # get the packet shell index
-        shells = sampled_packets_df.index.get_level_values(0)
+        shells = sampled_packets_df.index.get_level_values(1)
 
         # get the inner and outer velocity boundaries for each packet to compute
         # the initial radii
@@ -829,7 +829,7 @@ def create_packets(
 
         # packet_effective_times = sampled_packets_df.index.get_level_values(0)
 
-        decay_times = self.sample_decay_times(
+        times = self.sample_decay_times(
             isotopes, self.taus, self.effective_times, number_of_packets
         )
 
@@ -838,9 +838,7 @@ def create_packets(
         # )
         decay_time_indices = []
         for i in range(number_of_packets):
-            decay_time_indices.append(
-                get_index(decay_times[i], self.effective_times)
-            )
+            decay_time_indices.append(get_index(times[i], self.effective_times))
         # decay_times = self.create_packet_decay_times(
         #     number_of_packets, isotopes
         # )
@@ -879,7 +877,7 @@ def create_packets(
         nus_rf = np.zeros(number_of_packets)
 
         doppler_factors = doppler_factor_3D_all_packets(
-            directions, locations, decay_times
+            directions, locations, times
         )
 
         packet_energies_rf = packet_energies_cmf / doppler_factors
@@ -923,7 +921,7 @@ def create_packets(
             nus_cmf,
             statuses,
             shells,
-            decay_times,
+            times,
         )
 
     def calculate_positron_fraction(self, isotopes):
diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index c57ca720c29..95d3e872cd6 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/tardis/energy_input/main_gamma_ray_loop.py
@@ -331,6 +331,7 @@ def get_effective_time_array(time_start, time_end, time_space, time_steps):
 
 def run_gamma_ray_loop(
     model,
+    plasma,
     isotope_decay_df,
     cumulative_decays_df,
     num_decays,
@@ -370,17 +371,15 @@ def run_gamma_ray_loop(
             atom_number
         ].values
         if values.shape[0] > 1:
-            model.composition.isotopic_number_density.loc[
-                atom_number
-            ].update = np.sum(values, axis=0)
+            plasma.number_density.loc[atom_number].update = np.sum(
+                values, axis=0
+            )
         else:
-            model.composition.isotopic_number_density.loc[
-                atom_number
-            ].update = values
+            plasma.number_density.loc[atom_number].update = values
 
     # Electron number density
-    electron_number_density = model.composition.isotopic_number_density.mul(
-        model.composition.isotopic_number_density.index.get_level_values(0),
+    electron_number_density = plasma.number_density.mul(
+        plasma.number_density.index,
         axis=0,
     ).sum()
     electron_number = np.array(electron_number_density * ejecta_volume)
@@ -431,7 +430,7 @@ def run_gamma_ray_loop(
     )
     logger.info("Creating packets")
     packet_collection = packet_source.create_packets(
-        isotope_decay_df, num_decays
+        cumulative_decays_df, num_decays
     )
     logger.info("Creating packet list")
     packets = []
@@ -450,7 +449,7 @@ def run_gamma_ray_loop(
         for i in range(num_decays)
     ]
 
-    # return packets
+    return packets
 
     energy_bins = np.logspace(2, 3.8, spectrum_bins)
     energy_out = np.zeros((len(energy_bins - 1), time_steps))
@@ -491,7 +490,11 @@ def run_gamma_ray_loop(
     #     packets, times, inner_velocities, outer_velocities, effective_time_array
     # )
 
-    (energy_out, bin_width, packets_array,) = gamma_packet_loop(
+    (
+        energy_out,
+        bin_width,
+        packets_array,
+    ) = gamma_packet_loop(
         packets,
         grey_opacity,
         photoabsorption_opacity,

From 2089567a6a022699bba024a8f44d0a41547cff49 Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Mon, 24 Jun 2024 15:24:24 -0400
Subject: [PATCH 10/18] done some formatiing and changed what will come out as
 output

---
 tardis/energy_input/gamma_packet_loop.py   |  6 +-
 tardis/energy_input/main_gamma_ray_loop.py | 90 +++++++++++++---------
 2 files changed, 55 insertions(+), 41 deletions(-)

diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py
index 3100d82ae8b..e603e723a4f 100644
--- a/tardis/energy_input/gamma_packet_loop.py
+++ b/tardis/energy_input/gamma_packet_loop.py
@@ -212,7 +212,6 @@ def gamma_packet_loop(
                 effective_time_array[time_index],
                 effective_time_array[time_index + 1],
             )
-
             distance = min(
                 distance_interaction, distance_boundary, distance_time
             )
@@ -291,11 +290,10 @@ def gamma_packet_loop(
                         energy_bins[bin_index + 1] - energy_bins[bin_index]
                     )
                     freq_bin_width = bin_width / H_CGS_KEV
-                    energy_ergs = packet.energy_rf * KEV2ERG
-                    luminosity = energy_ergs / dt
                     energy_out[bin_index, time_index] += (
-                        energy_ergs / dt / freq_bin_width
+                        packet.energy_rf / dt / freq_bin_width
                     )
+                    luminosity = packet.energy_rf / dt
 
                     packet.status = GXPacketStatus.ESCAPED
                     escaped_packets += 1
diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index 95d3e872cd6..3eda7ba0d5e 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/tardis/energy_input/main_gamma_ray_loop.py
@@ -339,14 +339,60 @@ def run_gamma_ray_loop(
     effective_time_array,
     seed,
     positronium_fraction,
-    path_to_decay_data,
     spectrum_bins,
     time_steps,
     grey_opacity,
     photoabsorption_opacity="tardis",
     pair_creation_opacity="tardis",
 ):
+    """
+    Main loop to determine the gamma-ray propagation through the ejecta.
+
+    Parameters
+    ----------
+    model : tardis.model.Radial1DModel
+            Tardis model object.
+    plasma : tardis.plasma.standard_plasmas.BasePlasma
+            Tardis plasma object.
+    isotope_decay_df : pd.DataFrame
+            DataFrame containing the cumulative decay data.
+    cumulative_decays_df : pd.DataFrame
+            DataFrame containing the time evolving mass fractions.
+    num_decays : int
+            Number of packets to decay.
+    times : np.ndarray
+            Array of times in days.
+    effective_time_array : np.ndarray
+            Effective time array in days.
+    seed : int
+            Seed for the random number generator.
+    positronium_fraction : float
+            Fraction of positronium.
+    spectrum_bins : int
+            Number of spectrum bins.
+    time_steps : int
+            Number of time steps.
+    grey_opacity : float
+            Grey opacity.
+    photoabsorption_opacity : str
+            Photoabsorption opacity.
+    pair_creation_opacity : str
+            Pair creation opacity.
+
+
+    Returns
+    -------
+
+    escape_energy : pd.DataFrame
+            DataFrame containing the energy escaping the ejecta.
+    packets_df_escaped : pd.DataFrame
+            DataFrame containing the packets info that escaped the ejecta.
+
+
+    """
     np.random.seed(seed)
+    times = times * u.d.to(u.s)
+    effective_time_array = effective_time_array * u.d.to(u.s)
     inner_velocities = model.v_inner.to("cm/s").value
     outer_velocities = model.v_outer.to("cm/s").value
     ejecta_volume = model.volume.to("cm^3").value
@@ -390,25 +436,20 @@ def run_gamma_ray_loop(
     )
     mass_density_time = shell_masses[:, np.newaxis] * inv_volume_time
 
-    # gamma_ray_lines = get_nuclear_lines_database(path_to_decay_data)
     taus, parents = get_taus(raw_isotope_abundance)
     # Need to get the strings for the isotopes without the dashes
     taus = make_isotope_string_tardis_like(taus)
 
+    # Get only the gamma rays
     isotope_decay_df_gamma = isotope_decay_df[
         isotope_decay_df["radiation"] == "g"
     ]
+    total_energy_gamma = isotope_decay_df_gamma["decay_energy_erg"].sum()
 
-    total_energy_gamma = isotope_decay_df_gamma[
-        "decay_energy_keV"
-    ].sum()  # total energy in keV
-
-    total_energy_gamma_erg = total_energy_gamma * u.keV.to("erg")
-
-    energy_per_packet = total_energy_gamma_erg / num_decays
+    energy_per_packet = total_energy_gamma / num_decays
     energy_df_rows = np.zeros((number_of_shells, times.shape[0]))
 
-    logger.info(f"Total energy in gamma-rays is {total_energy_gamma_erg}")
+    logger.info(f"Total energy in gamma-rays is {total_energy_gamma}")
     logger.info(f"Energy per packet is {energy_per_packet}")
     average_power_per_mass = 1e41
     positron_energy_per_isotope = 1000  # keV
@@ -449,8 +490,6 @@ def run_gamma_ray_loop(
         for i in range(num_decays)
     ]
 
-    return packets
-
     energy_bins = np.logspace(2, 3.8, spectrum_bins)
     energy_out = np.zeros((len(energy_bins - 1), time_steps))
     packets_info_array = np.zeros((int(num_decays), 8))
@@ -465,34 +504,11 @@ def run_gamma_ray_loop(
         total_cmf_energy += p.energy_cmf
         total_rf_energy += p.energy_rf
 
-    print("Total CMF energy")
-    print(total_cmf_energy)
-    print("Total RF energy")
-    print(total_rf_energy)
-
-    # packets_df_escaped = pd.DataFrame(
-    #     data=packets_array,
-    #     columns=[
-    #         "packet_index",
-    #         "status",
-    #         "nu_cmf",
-    #         "nu_rf",
-    #         "energy_cmf",
-    #         "lum_rf",
-    #         "energy_rf",
-    #         "shell_number",
-    #     ],
-    # )
-
-    # return packets_df_escaped
-
-    # gamma_packet_loop_understanding(
-    #     packets, times, inner_velocities, outer_velocities, effective_time_array
-    # )
+    logger.info(f"Total CMF energy is {total_cmf_energy}")
+    logger.info(f"Total RF energy is {total_rf_energy}")
 
     (
         energy_out,
-        bin_width,
         packets_array,
     ) = gamma_packet_loop(
         packets,

From 4b9dba4b284023340a066637bbf3555113912a18 Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Mon, 24 Jun 2024 16:20:04 -0400
Subject: [PATCH 11/18] Cleanup

---
 tardis/energy_input/gamma_packet_loop.py      | 107 +------
 tardis/energy_input/gamma_ray_channel.py      |   3 +-
 tardis/energy_input/gamma_ray_grid.py         |  11 -
 .../energy_input/gamma_ray_packet_source.py   | 180 ++++--------
 tardis/energy_input/main_gamma_ray_loop.py    | 271 +-----------------
 tardis/energy_input/util.py                   |  20 +-
 6 files changed, 67 insertions(+), 525 deletions(-)

diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py
index e603e723a4f..5fd85b4824f 100644
--- a/tardis/energy_input/gamma_packet_loop.py
+++ b/tardis/energy_input/gamma_packet_loop.py
@@ -7,7 +7,6 @@
     compton_opacity_calculation,
     photoabsorption_opacity_calculation,
     pair_creation_opacity_calculation,
-    photoabsorption_opacity_calculation_kasen,
     kappa_calculation,
     pair_creation_opacity_artis,
     SIGMA_T,
@@ -15,15 +14,14 @@
 from tardis.energy_input.gamma_ray_grid import (
     distance_trace,
     move_packet,
-    calculate_distance_radial,
 )
 from tardis.energy_input.util import (
     doppler_factor_3d,
     C_CGS,
     H_CGS_KEV,
-    KEV2ERG,
     get_index,
 )
+
 from tardis.energy_input.GXPacket import GXPacketStatus
 from tardis.energy_input.gamma_ray_interactions import (
     get_compton_fraction_artis,
@@ -45,7 +43,6 @@ def gamma_packet_loop(
     pair_creation_opacity_type,
     electron_number_density_time,
     mass_density_time,
-    inv_volume_time,
     iron_group_fraction_per_shell,
     inner_velocities,
     outer_velocities,
@@ -109,23 +106,17 @@ def gamma_packet_loop(
     scattered_packets = 0
     packet_count = len(packets)
 
-    print("Packet count: ", packet_count)
-    # deposition_estimator = np.zeros_like(energy_df_rows)
+    print("Packet count:", packet_count)
 
     for i in range(packet_count):
         packet = packets[i]
         time_index = get_index(packet.time_current, times)
-        # print("Time index: ", time_index)
-        # print("Time current: ", packet.time_current)
-
         if time_index < 0:
             print(packet.time_current, time_index)
             raise ValueError("Packet time index less than 0!")
 
         scattered = False
 
-        initial_energy = packet.energy_cmf
-
         while packet.status == GXPacketStatus.IN_PROCESS:
             # Get delta-time value for this step
             dt = dt_array[time_index]
@@ -220,17 +211,6 @@ def gamma_packet_loop(
 
             packet = move_packet(packet, distance)
 
-            # deposition_estimator[packet.shell, time_index] += (
-            #    (initial_energy * 1000)
-            #    * distance
-            #    * (packet.energy_cmf / initial_energy)
-            #    * deposition_estimator_kasen(
-            #        comoving_energy,
-            #       mass_density_time[packet.shell, time_index],
-            #        iron_group_fraction_per_shell[packet.shell],
-            #    )
-            # )
-
             if distance == distance_time:
                 time_index += 1
 
@@ -252,27 +232,6 @@ def gamma_packet_loop(
 
                 packet, ejecta_energy_gained = process_packet_path(packet)
 
-                # Save packets to dataframe rows
-                # convert KeV to eV / s / cm^3
-                # energy_df_rows[packet.shell, time_index] += (
-                #    ejecta_energy_gained * 1000
-                # )
-
-                # energy_plot_df_rows[i] = np.array(
-                #    [
-                #        i,
-                #        ejecta_energy_gained * 1000
-                # * inv_volume_time[packet.shell, time_index]
-                #        / dt,
-                #        packet.get_location_r(),
-                #        packet.time_current,
-                #        packet.shell,
-                #        compton_opacity,
-                #        photoabsorption_opacity,
-                #        pair_creation_opacity,
-                #    ]
-                # )
-
                 if packet.status == GXPacketStatus.PHOTOABSORPTION:
                     # Packet destroyed, go to the next packet
                     break
@@ -322,7 +281,6 @@ def gamma_packet_loop(
 
     return (
         energy_out,
-        bin_width,
         packets_info_array,
     )
 
@@ -379,64 +337,3 @@ def process_packet_path(packet):
         ejecta_energy_gained = packet.energy_cmf
 
     return packet, ejecta_energy_gained
-
-
-@njit(**njit_dict_no_parallel)
-def gamma_packet_loop_understanding(
-    packets,
-    times,
-    inner_velocity,
-    outer_velocity,
-    effective_time_array,
-):
-
-    packet_count = len(packets)
-    print("Packet count: ", packet_count)
-    print("Effective time array: ", effective_time_array)
-    print("length of effective time array", len(effective_time_array))
-
-    for i in range(packet_count):
-        packet = packets[i]
-        time_index = get_index(packet.time_current, times)
-        print("Time index: ", time_index)
-        print("Time current: ", packet.time_current)
-        print("Time current in days: ", packet.time_current / 86400)
-        print("Length of time array:", len(times))
-
-        # we need to have the times within the simulation time steps
-
-        if time_index < 0:
-            print(packet.time_current, time_index)
-            raise ValueError("Packet time index less than 0!")
-        elif time_index >= len(times):
-            print(packet.time_current, time_index)
-            raise ValueError("Packet time index greater than time steps!")
-
-        # while packet.status == GXPacketStatus.IN_PROCESS:
-
-        print(
-            "Inner shell velocity: ",
-            inner_velocity[packet.shell],
-        )
-        print(
-            "Outer shell velocity: ",
-            outer_velocity[packet.shell],
-        )
-
-        print(
-            "Effective time array: ",
-            effective_time_array[time_index],
-        )
-        print(
-            "Effective time array in days:",
-            effective_time_array[time_index] / 86400,
-        )
-
-        distance_boundary, shell_change = calculate_distance_radial(
-            packet,
-            inner_velocity[packet.shell] * effective_time_array[time_index],
-            outer_velocity[packet.shell] * effective_time_array[time_index],
-        )
-
-        print("Distance boundary: ", distance_boundary)
-        print("Shell change: ", shell_change)
diff --git a/tardis/energy_input/gamma_ray_channel.py b/tardis/energy_input/gamma_ray_channel.py
index 7321923d90e..7b0f480a608 100644
--- a/tardis/energy_input/gamma_ray_channel.py
+++ b/tardis/energy_input/gamma_ray_channel.py
@@ -110,7 +110,8 @@ def calculate_total_decays(inventories, time_delta):
 
 def create_isotope_decay_df(cumulative_decay_df, gamma_ray_lines):
     """
-    Function to create a dataframe of isotopes for each shell with their decay mode, number of decays, radiation type,
+    Function to create a dataframe of isotopes for each shell with their decay mode,
+    number of decays, radiation type,
     radiation energy and radiation intensity.
 
     Parameters
diff --git a/tardis/energy_input/gamma_ray_grid.py b/tardis/energy_input/gamma_ray_grid.py
index 46a27cc3971..711fdaffafb 100644
--- a/tardis/energy_input/gamma_ray_grid.py
+++ b/tardis/energy_input/gamma_ray_grid.py
@@ -50,17 +50,6 @@ def calculate_distance_radial(photon, r_inner, r_outer):
         outer_2 = -1
 
     distances = np.array([inner_1, inner_2, outer_1, outer_2])
-
-    # print(
-    #     "Distance to inner shell boundaries: ",
-    #     photon.get_location_r() - r_inner,
-    # )
-    # print(
-    #     "Distance to outer shell boundaries: ",
-    #     photon.get_location_r() - r_outer,
-    # )
-
-    # the correct distance is the shortest positive distance
     distance_list = [i for i in distances if i > 0]
 
     if not distance_list:
diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index 8fb1ecc8927..5096850304a 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -15,7 +15,6 @@
     doppler_factor_3D_all_packets,
     get_index,
     get_random_unit_vector,
-    C_CGS,
 )
 from tardis.transport.montecarlo.packet_source import BasePacketSource
 
@@ -438,18 +437,18 @@ def create_packets(self, decays_per_isotope, *args, **kwargs):
                 packet_energies_cmf[
                     packet_index - isotope_packet_count : packet_index
                 ] = initial_packet_energies_cmf
-                nus_rf[packet_index - isotope_packet_count : packet_index] = (
-                    initial_nus_rf
-                )
-                nus_cmf[packet_index - isotope_packet_count : packet_index] = (
-                    initial_nus_cmf
-                )
-                shells[packet_index - isotope_packet_count : packet_index] = (
-                    initial_shells
-                )
-                times[packet_index - isotope_packet_count : packet_index] = (
-                    initial_times
-                )
+                nus_rf[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_nus_rf
+                nus_cmf[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_nus_cmf
+                shells[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_shells
+                times[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_times
 
         return GXPacketCollection(
             locations,
@@ -496,14 +495,42 @@ def __init__(
         outer_velocities,
         inv_volume_time,
         times,
-        energy_df_rows,
         effective_times,
         taus,
         parents,
-        average_positron_energies,
-        average_power_per_mass,
         **kwargs,
     ):
+        """
+        New Gamma ray packet source class
+
+        Parameters
+        ----------
+
+        packet_energy : float
+            Energy of the gamma ray packet
+        isotope_decay_df : pd.DataFrame
+            DataFrame of isotope decay data
+        positronium_fraction : float
+            Fraction of positrons that form positronium
+        inner_velocities : array
+            Array of inner shell velocities
+        outer_velocities : array
+            Array of outer shell velocities
+        inv_volume_time : array
+            Array of inverse volume times (please explain)
+        times : array
+            Array of time steps
+        effective_times : array
+            Array of effective time steps
+        taus : dict
+            Dictionary of isotope mean lifetimes in seconds
+        parents : dict
+            Dictionary of isotope parents
+        average_positron_energies : dict
+            Dictionary of average positron energies per isotope
+
+
+        """
         self.packet_energy = packet_energy
         self.isotope_decay_df = isotope_decay_df
         self.positronium_fraction = positronium_fraction
@@ -511,13 +538,9 @@ def __init__(
         self.outer_velocities = outer_velocities
         self.inv_volume_time = inv_volume_time
         self.times = times
-        self.energy_df_rows = energy_df_rows
         self.effective_times = effective_times
         self.taus = taus
         self.parents = parents
-        self.average_positron_energies = average_positron_energies
-        self.average_power_per_mass = average_power_per_mass
-        self.energy_plot_positron_rows = np.empty(0)
         super().__init__(**kwargs)
 
     def create_packet_mus(self, no_of_packets, *args, **kwargs):
@@ -651,7 +674,9 @@ def create_packet_times_uniform_time(self, no_of_packets, start, end):
     def create_packet_times_uniform_energy(
         self, no_of_packets, isotopes, decay_time
     ):
-        """Samples the decay time from the mean lifetime of the isotopes
+        """
+        This method is not used in the current implementation
+        Samples the decay time from the mean lifetime of the isotopes
 
         Parameters
         ----------
@@ -683,30 +708,6 @@ def create_packet_times_uniform_energy(
                 )
         return decay_times
 
-    def create_packet_decay_times(self, number_of_packets, isotopes):
-        """Samples the decay time from the mean lifetime of the isotopes
-
-        Parameters
-        ----------
-        number_of_packets : int
-            Number of packets
-        times : array
-            Array of time steps
-
-        Returns
-        -------
-        array
-            Array of decay times
-        """
-        decay_times = np.zeros(number_of_packets)
-        for i, isotope in enumerate(isotopes.to_numpy()):
-            if isotope in self.taus:
-                decay_times[i] = -self.taus[isotope] * np.log(
-                    np.random.random()
-                )
-
-        return decay_times
-
     def sample_decay_times(self, isotopes, taus, times, number_of_packets):
         """
         Sample decay times for a given number of packets
@@ -741,7 +742,6 @@ def sample_decay_times(self, isotopes, taus, times, number_of_packets):
             while (decay_times[i] <= decay_time_min) or (
                 decay_times[i] >= decay_time_max
             ):
-                # since we know which isotope the packet is associated with, we can use the tau value for that isotope
                 decay_times[i] = -taus[isotope] * np.log(np.random.random())
 
         return decay_times
@@ -771,7 +771,6 @@ def create_packets(
         nus_rf = np.zeros(number_of_packets)
         nus_cmf = np.zeros(number_of_packets)
         times = np.zeros(number_of_packets)
-        # set packets to IN_PROCESS status
         statuses = np.ones(number_of_packets, dtype=np.int64) * 3
 
         self.energy_plot_positron_rows = np.zeros((number_of_packets, 4))
@@ -779,9 +778,7 @@ def create_packets(
         # compute positronium continuum
         positronium_energy, positronium_intensity = positronium_continuum()
 
-        # sample packets from dataframe, returning a dataframe where each row is
-        # a sampled packet
-
+        # sample packets from the gamma-ray lines
         sampled_packets_df_gamma = decays_per_isotope[
             decays_per_isotope["radiation"] == "g"
         ]
@@ -791,13 +788,8 @@ def create_packets(
             replace=True,
             random_state=np.random.RandomState(self.base_seed),
         )
-        # get unique isotopes that have produced packets
+        # get the isotopes and shells of the sampled packets
         isotopes = sampled_packets_df.index.get_level_values(2)
-
-        # compute the positron fraction for unique isotopes
-        isotope_positron_fraction = self.calculate_positron_fraction(isotopes)
-
-        # get the packet shell index
         shells = sampled_packets_df.index.get_level_values(1)
 
         # get the inner and outer velocity boundaries for each packet to compute
@@ -807,42 +799,16 @@ def create_packets(
 
         # sample radii at time = 0
         initial_radii = self.create_packet_radii(sampled_packets_df)
-
-        # get the time step index of the packets
-        # Comment: Need to discuss. the indices should be used
-        # We can use the indices of the effective time array used as a timestamp for each packet
-
-        # packet_effective_times = sampled_packets_df.index.get_level_values(0)
-        # initial_time_indexes = np.indices(
-        #    sampled_packets_df.index.get_level_values(0).shape
-        # )
-
-        # get the time of the middle of the step for each packet
-        # Comment: Array indices need to be integers
-        # times_indices = np.indices(times.shape).flatten()
-
-        # scale radius by packet decay time. This could be replaced with
-        # Geometry object calculations. Note that this also adds a random
-        # unit vector multiplication for 3D. May not be needed.
-
-        # Comment: initial radii needs to be an array
-
-        # packet_effective_times = sampled_packets_df.index.get_level_values(0)
-
+        # sample decay times
         times = self.sample_decay_times(
             isotopes, self.taus, self.effective_times, number_of_packets
         )
 
-        # decay_time_indices = np.digitize(
-        #    decay_times, sorted(packet_effective_times.values), right=False
-        # )
+        # get the time step index of the packets
         decay_time_indices = []
         for i in range(number_of_packets):
             decay_time_indices.append(get_index(times[i], self.effective_times))
-        # decay_times = self.create_packet_decay_times(
-        #     number_of_packets, isotopes
-        # )
-
+        # 3D locations
         locations = (
             initial_radii.values
             * self.effective_times[decay_time_indices]
@@ -852,8 +818,6 @@ def create_packets(
         # sample directions (valid at all times), non-relativistic
         directions = self.create_packet_directions(number_of_packets)
 
-        # the individual gamma-ray energy that makes up a packet
-        # co-moving frame, including positronium formation
         nu_energies_cmf = self.create_packet_nus(
             number_of_packets,
             sampled_packets_df,
@@ -862,17 +826,12 @@ def create_packets(
             positronium_intensity,
         )
 
-        # equivalent frequencies
         nus_cmf = nu_energies_cmf / H_CGS_KEV
 
-        # per packet co-moving frame total energy
         packet_energies_cmf = self.create_packet_energies(
             number_of_packets, self.packet_energy
         )
 
-        # rest frame gamma-ray energy and frequency
-        # this probably works fine without the loop
-        # non-relativistic
         packet_energies_rf = np.zeros(number_of_packets)
         nus_rf = np.zeros(number_of_packets)
 
@@ -882,35 +841,6 @@ def create_packets(
 
         packet_energies_rf = packet_energies_cmf / doppler_factors
         nus_rf = nus_cmf / doppler_factors
-        # The for loop makes this program slow. Changed to numpy multiply
-
-        # deposit positron energy in both output arrays
-        # this is an average across all packets that are created
-        # it could be changed to be only for packets that are from positrons
-
-        # self.energy_plot_positron_rows[i] = np.array(
-        #     [
-        #         i,
-        #         isotope_positron_fraction[
-        #             sampled_packets_df.index.get_level_values(2)[i]
-        #         ]
-        #         * packet_energies_cmf[i],
-        #         # this needs to be sqrt(sum of squares) to get radius
-        #         np.linalg.norm(locations[i]),
-        #         times[i],
-        #     ]
-        # )
-
-        # this is an average across all packets that are created
-        # it could be changed to be only for packets that are from positrons
-
-        # comment: we need the indices here
-        # self.energy_df_rows[shells[i], times_indices[i]] += (
-        #     isotope_positron_fraction[
-        #         sampled_packets_df.index.get_level_values(2)[i]
-        #     ]
-        #     * packet_energies_cmf[i]
-        # )
 
         return GXPacketCollection(
             locations,
@@ -960,11 +890,3 @@ def calculate_positron_fraction(self, isotopes):
                 / total_energy_per_isotope[isotope]
             )
         return positron_fraction
-
-        # for isotope in isotopes:
-        #    isotope_energy = self.gamma_ray_lines[isotope][0, :]
-        #    isotope_intensity = self.gamma_ray_lines[isotope][1, :]
-        #    positron_fraction[isotope] = self.average_positron_energies[
-        #        isotope
-        #    ] / np.sum(isotope_energy * isotope_intensity)
-        # return positron_fraction
diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index 3eda7ba0d5e..0f045e61f8e 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/tardis/energy_input/main_gamma_ray_loop.py
@@ -7,15 +7,8 @@
 from tardis.energy_input.energy_source import (
     get_nuclear_lines_database,
 )
-from tardis.energy_input.gamma_packet_loop import (
-    gamma_packet_loop,
-    gamma_packet_loop_understanding,
-)
-from tardis.energy_input.gamma_ray_channel import (
-    calculate_total_decays,
-    create_inventories_dict,
-    create_isotope_dicts,
-)
+from tardis.energy_input.gamma_packet_loop import gamma_packet_loop
+
 
 from tardis.energy_input.gamma_ray_transport import (
     calculate_total_decays_old,
@@ -43,258 +36,6 @@
 logging.basicConfig(level=logging.INFO)
 
 
-def run_gamma_ray_loop_old(
-    model,
-    plasma,
-    num_decays,
-    time_start,
-    time_end,
-    time_space,
-    time_steps,
-    seed,
-    positronium_fraction,
-    path_to_decay_data,
-    spectrum_bins,
-    grey_opacity,
-    photoabsorption_opacity="tardis",
-    pair_creation_opacity="tardis",
-):
-    """
-    Main loop to determine the gamma-ray propagation through the ejecta.
-    """
-    np.random.seed(seed)
-    time_explosion = model.time_explosion.to(u.s).value
-    inner_velocities = model.v_inner.to("cm/s").value
-    outer_velocities = model.v_outer.to("cm/s").value
-    ejecta_volume = model.volume.to("cm^3").value
-    number_of_shells = model.no_of_shells
-    shell_masses = model.volume * model.density
-    raw_isotope_abundance = model.composition.raw_isotope_abundance.sort_values(
-        by=["atomic_number", "mass_number"], ascending=False
-    )
-    time_start *= u.d.to(u.s)
-    time_end *= u.d.to(u.s)
-
-    assert time_start < time_end, "time_start must be smaller than time_end!"
-    if time_space == "log":
-        times = np.geomspace(time_start, time_end, time_steps + 1)
-        effective_time_array = np.sqrt(times[:-1] * times[1:])
-    else:
-        times = np.linspace(time_start, time_end, time_steps + 1)
-        effective_time_array = 0.5 * (times[:-1] + times[1:])
-
-    dt_array = np.diff(times)
-
-    ejecta_velocity_volume = calculate_ejecta_velocity_volume(model)
-
-    inv_volume_time = (
-        1.0 / ejecta_velocity_volume[:, np.newaxis]
-    ) / effective_time_array**3.0
-
-    energy_df_rows = np.zeros((number_of_shells, time_steps))
-    # Use isotopic number density
-    for atom_number in plasma.isotope_number_density.index.get_level_values(0):
-        values = plasma.isotope_number_density.loc[atom_number].values
-        if values.shape[0] > 1:
-            plasma.isotope_number_density.loc[atom_number].update = np.sum(
-                values, axis=0
-            )
-        else:
-            plasma.isotope_number_density.loc[atom_number].update = values
-
-    # Electron number density
-    electron_number_density = plasma.number_density.mul(
-        plasma.number_density.index, axis=0
-    ).sum()
-    taus, parents = get_taus(raw_isotope_abundance)
-    # inventories = raw_isotope_abundance.to_inventories()
-    electron_number = np.array(electron_number_density * ejecta_volume)
-    electron_number_density_time = (
-        electron_number[:, np.newaxis] * inv_volume_time
-    )
-
-    # Calculate decay chain energies
-
-    mass_density_time = shell_masses[:, np.newaxis] * inv_volume_time
-    gamma_ray_lines = get_nuclear_lines_database(path_to_decay_data)
-    isotope_dict = create_isotope_dicts_old(raw_isotope_abundance, shell_masses)
-    inventories_dict = create_inventories_dict_old(isotope_dict)
-    total_decays = calculate_total_decays_old(
-        inventories_dict, time_end - time_start
-    )
-
-    (
-        average_energies,
-        average_positron_energies,
-        gamma_ray_line_dict,
-    ) = calculate_average_energies(raw_isotope_abundance, gamma_ray_lines)
-
-    decayed_energy = decay_chain_energies(
-        average_energies,
-        total_decays,
-    )
-    energy_per_mass, energy_df = calculate_energy_per_mass(
-        decayed_energy, raw_isotope_abundance, shell_masses
-    )
-    average_power_per_mass = calculate_average_power_per_mass(
-        energy_per_mass, time_end - time_start
-    )
-    number_of_isotopes = plasma.isotope_number_density * ejecta_volume
-    total_isotope_number = number_of_isotopes.sum().sum()
-    decayed_packet_count = num_decays * number_of_isotopes.divide(
-        total_isotope_number, axis=0
-    )
-
-    total_energy = energy_df.sum().sum()
-    energy_per_packet = total_energy / num_decays
-    packets_per_isotope_df = (
-        distribute_packets(decayed_energy, total_energy, num_decays)
-        .round()
-        .fillna(0)
-        .astype(int)
-    )
-
-    total_energy = total_energy * u.eV.to("erg")
-
-    logger.info(f"Total gamma-ray energy is {total_energy}")
-
-    iron_group_fraction = iron_group_fraction_per_shell(model)
-    number_of_packets = packets_per_isotope_df.sum().sum()
-    logger.info(f"Total number of packets is {number_of_packets}")
-    individual_packet_energy = total_energy / number_of_packets
-    logger.info(f"Energy per packet is {individual_packet_energy}")
-
-    logger.info("Initializing packets")
-
-    packet_source = RadioactivePacketSource(
-        individual_packet_energy,
-        gamma_ray_line_dict,
-        positronium_fraction,
-        inner_velocities,
-        outer_velocities,
-        inv_volume_time,
-        times,
-        energy_df_rows,
-        effective_time_array,
-        taus,
-        parents,
-        average_positron_energies,
-        average_power_per_mass,
-    )
-
-    packet_collection = packet_source.create_packets(packets_per_isotope_df)
-
-    energy_df_rows = packet_source.energy_df_rows
-    energy_plot_df_rows = np.zeros((number_of_packets, 8))
-
-    logger.info("Creating packet list")
-    packets = []
-    total_cmf_energy = packet_collection.energy_cmf.sum()
-    total_rf_energy = packet_collection.energy_rf.sum()
-    for i in range(number_of_packets):
-        packet = GXPacket(
-            packet_collection.location[:, i],
-            packet_collection.direction[:, i],
-            packet_collection.energy_rf[i],
-            packet_collection.energy_cmf[i],
-            packet_collection.nu_rf[i],
-            packet_collection.nu_cmf[i],
-            packet_collection.status[i],
-            packet_collection.shell[i],
-            packet_collection.time_current[i],
-        )
-        packets.append(packet)
-        energy_plot_df_rows[i] = np.array(
-            [
-                i,
-                packet.energy_rf,
-                packet.get_location_r(),
-                packet.time_current,
-                int(packet.status),
-                0,
-                0,
-                0,
-            ]
-        )
-
-    logger.info(f"Total cmf energy is {total_cmf_energy}")
-    logger.info(f"Total rf energy is {total_rf_energy}")
-
-    energy_bins = np.logspace(2, 3.8, spectrum_bins)
-    energy_out = np.zeros((len(energy_bins - 1), time_steps))
-    packets_info_array = np.zeros((int(num_decays), 8))
-
-    (
-        energy_df_rows,
-        energy_plot_df_rows,
-        energy_out,
-        deposition_estimator,
-        bin_width,
-        packets_array,
-    ) = gamma_packet_loop(
-        packets,
-        grey_opacity,
-        photoabsorption_opacity,
-        pair_creation_opacity,
-        electron_number_density_time,
-        mass_density_time,
-        inv_volume_time,
-        iron_group_fraction.to_numpy(),
-        inner_velocities,
-        outer_velocities,
-        times,
-        dt_array,
-        effective_time_array,
-        energy_bins,
-        energy_df_rows,
-        energy_plot_df_rows,
-        energy_out,
-        packets_info_array,
-    )
-
-    energy_plot_df = pd.DataFrame(
-        data=energy_plot_df_rows,
-        columns=[
-            "packet_index",
-            "energy_input",
-            "energy_input_r",
-            "energy_input_time",
-            "energy_input_type",
-            "compton_opacity",
-            "photoabsorption_opacity",
-            "total_opacity",
-        ],
-    )
-
-    energy_plot_positrons = pd.DataFrame(
-        data=packet_source.energy_plot_positron_rows,
-        columns=[
-            "packet_index",
-            "energy_input",
-            "energy_input_r",
-            "energy_input_time",
-        ],
-    )
-
-    energy_estimated_deposition = (
-        pd.DataFrame(data=deposition_estimator, columns=times[:-1])
-    ) / dt_array
-
-    energy_df = pd.DataFrame(data=energy_df_rows, columns=times[:-1]) / dt_array
-    escape_energy = pd.DataFrame(
-        data=energy_out, columns=times[:-1], index=energy_bins
-    )
-
-    return (
-        energy_df,
-        energy_plot_df,
-        escape_energy,
-        decayed_packet_count,
-        energy_plot_positrons,
-        energy_estimated_deposition,
-    )
-
-
 def get_effective_time_array(time_start, time_end, time_space, time_steps):
     """
     Function to get the effective time array for the gamma-ray loop.
@@ -396,7 +137,6 @@ def run_gamma_ray_loop(
     inner_velocities = model.v_inner.to("cm/s").value
     outer_velocities = model.v_outer.to("cm/s").value
     ejecta_volume = model.volume.to("cm^3").value
-    number_of_shells = model.no_of_shells
     shell_masses = model.volume * model.density
     raw_isotope_abundance = model.composition.raw_isotope_abundance.sort_values(
         by=["atomic_number", "mass_number"], ascending=False
@@ -447,12 +187,9 @@ def run_gamma_ray_loop(
     total_energy_gamma = isotope_decay_df_gamma["decay_energy_erg"].sum()
 
     energy_per_packet = total_energy_gamma / num_decays
-    energy_df_rows = np.zeros((number_of_shells, times.shape[0]))
 
     logger.info(f"Total energy in gamma-rays is {total_energy_gamma}")
     logger.info(f"Energy per packet is {energy_per_packet}")
-    average_power_per_mass = 1e41
-    positron_energy_per_isotope = 1000  # keV
 
     packet_source = GammaRayPacketSource(
         energy_per_packet,
@@ -462,12 +199,9 @@ def run_gamma_ray_loop(
         outer_velocities,
         inv_volume_time,
         times,
-        energy_df_rows,
         effective_time_array,
         taus,
         parents,
-        positron_energy_per_isotope,
-        average_power_per_mass,
     )
     logger.info("Creating packets")
     packet_collection = packet_source.create_packets(
@@ -517,7 +251,6 @@ def run_gamma_ray_loop(
         pair_creation_opacity,
         electron_number_density_time,
         mass_density_time,
-        inv_volume_time,
         iron_group_fraction.to_numpy(),
         inner_velocities,
         outer_velocities,
diff --git a/tardis/energy_input/util.py b/tardis/energy_input/util.py
index 6fb8eb7a24a..ebbf17298e5 100644
--- a/tardis/energy_input/util.py
+++ b/tardis/energy_input/util.py
@@ -186,13 +186,13 @@ def solve_quadratic_equation(position, direction, radius):
     a = np.sum(direction**2)
     b = 2.0 * np.sum(position * direction)
     c = -(radius**2) + np.sum(position**2)
-    root = b**2 - 4 * a * c
+    discriminant = b**2 - 4 * a * c
     solution_1 = -np.inf
     solution_2 = -np.inf
-    if root > 0.0:
-        solution_1 = (-b + np.sqrt(root)) / (2 * a)
-        solution_2 = (-b - np.sqrt(root)) / (2 * a)
-    elif root == 0:
+    if discriminant > 0.0:
+        solution_1 = (-b + np.sqrt(discriminant)) / (2 * a)
+        solution_2 = (-b - np.sqrt(discriminant)) / (2 * a)
+    elif discriminant == 0:
         solution_1 = -b / (2 * a)
 
     return solution_1, solution_2
@@ -220,13 +220,13 @@ def solve_quadratic_equation_expanding(position, direction, time, radius):
     a = np.dot(direction, direction) - (radius / light_distance) ** 2.0
     b = 2.0 * (np.dot(position, direction) - radius**2.0 / light_distance)
     c = np.dot(position, position) - radius**2.0
-    root = b**2.0 - 4.0 * a * c
+    discriminant = b**2.0 - 4.0 * a * c
     solution_1 = -np.inf
     solution_2 = -np.inf
-    if root > 0.0:
-        solution_1 = (-b + np.sqrt(root)) / (2.0 * a)
-        solution_2 = (-b - np.sqrt(root)) / (2.0 * a)
-    elif root == 0:
+    if discriminant > 0.0:
+        solution_1 = (-b + np.sqrt(discriminant)) / (2.0 * a)
+        solution_2 = (-b - np.sqrt(discriminant)) / (2.0 * a)
+    elif discriminant == 0:
         solution_1 = -b / (2.0 * a)
 
     return solution_1, solution_2

From f24a2c9ef27c429ba025382f5cae358046d032e8 Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Mon, 24 Jun 2024 16:50:09 -0400
Subject: [PATCH 12/18] Fixed black issues

---
 tardis/energy_input/main_gamma_ray_loop.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index 0f045e61f8e..31ee2c447a3 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/tardis/energy_input/main_gamma_ray_loop.py
@@ -241,10 +241,7 @@ def run_gamma_ray_loop(
     logger.info(f"Total CMF energy is {total_cmf_energy}")
     logger.info(f"Total RF energy is {total_rf_energy}")
 
-    (
-        energy_out,
-        packets_array,
-    ) = gamma_packet_loop(
+    (energy_out, packets_array,) = gamma_packet_loop(
         packets,
         grey_opacity,
         photoabsorption_opacity,

From 12a72f486acc9eb028dd88cc844541376cc5f96b Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Fri, 28 Jun 2024 11:50:29 -0400
Subject: [PATCH 13/18] Changed sampling of time to uniform time

---
 .../energy_input/gamma_ray_packet_source.py   | 57 +++++++++++++------
 1 file changed, 39 insertions(+), 18 deletions(-)

diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index 5096850304a..85213bd6813 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -437,18 +437,18 @@ def create_packets(self, decays_per_isotope, *args, **kwargs):
                 packet_energies_cmf[
                     packet_index - isotope_packet_count : packet_index
                 ] = initial_packet_energies_cmf
-                nus_rf[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_nus_rf
-                nus_cmf[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_nus_cmf
-                shells[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_shells
-                times[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_times
+                nus_rf[packet_index - isotope_packet_count : packet_index] = (
+                    initial_nus_rf
+                )
+                nus_cmf[packet_index - isotope_packet_count : packet_index] = (
+                    initial_nus_cmf
+                )
+                shells[packet_index - isotope_packet_count : packet_index] = (
+                    initial_shells
+                )
+                times[packet_index - isotope_packet_count : packet_index] = (
+                    initial_times
+                )
 
         return GXPacketCollection(
             locations,
@@ -746,6 +746,19 @@ def sample_decay_times(self, isotopes, taus, times, number_of_packets):
 
         return decay_times
 
+    def sample_decay_times_uniformly(
+        self, times, time_indices, number_of_packets
+    ):
+
+        decay_times = np.zeros(number_of_packets)
+
+        for i in range(number_of_packets):
+            decay_times[i] = np.random.uniform(
+                times[time_indices[i]], times[time_indices[i] + 1]
+            )
+
+        return decay_times
+
     def create_packets(
         self, decays_per_isotope, number_of_packets, *args, **kwargs
     ):
@@ -800,14 +813,22 @@ def create_packets(
         # sample radii at time = 0
         initial_radii = self.create_packet_radii(sampled_packets_df)
         # sample decay times
-        times = self.sample_decay_times(
-            isotopes, self.taus, self.effective_times, number_of_packets
+        sampled_times = sampled_packets_df.index.get_level_values("time")
+        sampled_time_indices = np.searchsorted(times, sampled_times)
+
+        decay_times = self.sample_decay_times_uniformly(
+            sampled_times, sampled_time_indices, number_of_packets
         )
 
         # get the time step index of the packets
-        decay_time_indices = []
-        for i in range(number_of_packets):
-            decay_time_indices.append(get_index(times[i], self.effective_times))
+        # decay_time_indices = []
+        # for i in range(number_of_packets):
+        #    decay_time_indices.append(
+        #        get_index(decay_times[i], self.effective_times)
+        #    )
+
+        # get the time of the middle of the step for each packet
+        decay_time_indices = np.searchsorted(self.effective_times, decay_times)
         # 3D locations
         locations = (
             initial_radii.values
@@ -851,7 +872,7 @@ def create_packets(
             nus_cmf,
             statuses,
             shells,
-            times,
+            decay_times,
         )
 
     def calculate_positron_fraction(self, isotopes):

From 2215b88da893459e4c5b41e3dad6b03f05211fc0 Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Mon, 1 Jul 2024 16:12:38 -0400
Subject: [PATCH 14/18] Uniform time sampling between adjacent time steps

---
 tardis/energy_input/gamma_packet_loop.py      |  4 +-
 .../energy_input/gamma_ray_packet_source.py   | 58 ++++++++++++-------
 tardis/energy_input/main_gamma_ray_loop.py    |  9 ++-
 3 files changed, 46 insertions(+), 25 deletions(-)

diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py
index bb5906354c2..24708b6200e 100644
--- a/tardis/energy_input/gamma_packet_loop.py
+++ b/tardis/energy_input/gamma_packet_loop.py
@@ -3,7 +3,7 @@
 from numba import njit
 
 from tardis.transport.montecarlo import njit_dict_no_parallel
-from tardis.transport.montecarlo.opacities import (
+from tardis.opacities.opacities import (
     compton_opacity_calculation,
     photoabsorption_opacity_calculation,
     pair_creation_opacity_calculation,
@@ -114,7 +114,7 @@ def gamma_packet_loop(
 
     for i in range(packet_count):
         packet = packets[i]
-        time_index = get_index(packet.time_current, times)
+        time_index = get_index(packet.time_current, effective_time_array)
         if time_index < 0:
             print(packet.time_current, time_index)
             raise ValueError("Packet time index less than 0!")
diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index 85213bd6813..fab361a234f 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -746,21 +746,37 @@ def sample_decay_times(self, isotopes, taus, times, number_of_packets):
 
         return decay_times
 
-    def sample_decay_times_uniformly(
-        self, times, time_indices, number_of_packets
-    ):
+    def sample_decay_times_uniformly(self, time, number_of_packets, seed):
+        """
+        Sample decay times uniformly
+
+        Parameters
+        ----------
+        time : array
+            Array of time steps in seconds
+        number_of_packets : int
+            Number of packets
+
+        Returns
+        -------
+        decay_times : array
+            Array of decay times for each packet
+
+        """
 
         decay_times = np.zeros(number_of_packets)
 
-        for i in range(number_of_packets):
-            decay_times[i] = np.random.uniform(
-                times[time_indices[i]], times[time_indices[i] + 1]
-            )
+        np.random.seed(seed)
+        decay_times[0] = self.times[0]
+        decay_times[-1] = self.times[-1]
+
+        for j in range(1, number_of_packets - 1):
+            decay_times[j] = np.random.uniform(time[j - 1], time[j])
 
         return decay_times
 
     def create_packets(
-        self, decays_per_isotope, number_of_packets, *args, **kwargs
+        self, decays_per_isotope, number_of_packets, seed, *args, **kwargs
     ):
         """Initialize a collection of GXPacket objects for the simulation
         to operate on.
@@ -813,23 +829,25 @@ def create_packets(
         # sample radii at time = 0
         initial_radii = self.create_packet_radii(sampled_packets_df)
         # sample decay times
-        sampled_times = sampled_packets_df.index.get_level_values("time")
-        sampled_time_indices = np.searchsorted(times, sampled_times)
+        sampled_times = (
+            sampled_packets_df.index.get_level_values("time") * 86400.0
+        )
+        # Get the indices of the time steps from the sampled times
 
         decay_times = self.sample_decay_times_uniformly(
-            sampled_times, sampled_time_indices, number_of_packets
+            sampled_times, number_of_packets, seed
         )
 
-        # get the time step index of the packets
-        # decay_time_indices = []
-        # for i in range(number_of_packets):
-        #    decay_time_indices.append(
-        #        get_index(decay_times[i], self.effective_times)
-        #    )
-
         # get the time of the middle of the step for each packet
-        decay_time_indices = np.searchsorted(self.effective_times, decay_times)
+        # decay_time_indices = np.searchsorted(self.effective_times, decay_times)
         # 3D locations
+
+        decay_time_indices = []
+        for i in range(number_of_packets):
+            decay_time_indices.append(
+                get_index(decay_times[i], self.effective_times)
+            )
+
         locations = (
             initial_radii.values
             * self.effective_times[decay_time_indices]
@@ -857,7 +875,7 @@ def create_packets(
         nus_rf = np.zeros(number_of_packets)
 
         doppler_factors = doppler_factor_3D_all_packets(
-            directions, locations, times
+            directions, locations, decay_times
         )
 
         packet_energies_rf = packet_energies_cmf / doppler_factors
diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index 31ee2c447a3..6f8fad0bcac 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/tardis/energy_input/main_gamma_ray_loop.py
@@ -205,7 +205,7 @@ def run_gamma_ray_loop(
     )
     logger.info("Creating packets")
     packet_collection = packet_source.create_packets(
-        cumulative_decays_df, num_decays
+        cumulative_decays_df, num_decays, seed
     )
     logger.info("Creating packet list")
     packets = []
@@ -241,7 +241,10 @@ def run_gamma_ray_loop(
     logger.info(f"Total CMF energy is {total_cmf_energy}")
     logger.info(f"Total RF energy is {total_rf_energy}")
 
-    (energy_out, packets_array,) = gamma_packet_loop(
+    (
+        energy_out,
+        packets_array,
+    ) = gamma_packet_loop(
         packets,
         grey_opacity,
         photoabsorption_opacity,
@@ -274,7 +277,7 @@ def run_gamma_ray_loop(
     )
 
     escape_energy = pd.DataFrame(
-        data=energy_out, columns=times[:-1], index=energy_bins
+        data=energy_out, columns=effective_time_array, index=energy_bins
     )
 
     return escape_energy, packets_df_escaped

From 4c027f54f9cf8426327b81f77ebd0cb444ac4232 Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Tue, 16 Jul 2024 09:24:47 -0400
Subject: [PATCH 15/18] corrected the total number of decays

---
 tardis/energy_input/gamma_ray_channel.py | 40 ++++++++++--------------
 1 file changed, 17 insertions(+), 23 deletions(-)

diff --git a/tardis/energy_input/gamma_ray_channel.py b/tardis/energy_input/gamma_ray_channel.py
index 7b0f480a608..d565c61ce63 100644
--- a/tardis/energy_input/gamma_ray_channel.py
+++ b/tardis/energy_input/gamma_ray_channel.py
@@ -171,7 +171,7 @@ def create_isotope_decay_df(cumulative_decay_df, gamma_ray_lines):
     return isotope_decay_df
 
 
-def evolve_mass_fraction(raw_isotope_mass_fraction, time_array):
+def time_evolve_mass_fraction(raw_isotope_mass_fraction, time_array):
     """
     Function to evolve the mass fraction of isotopes with time.
 
@@ -206,7 +206,7 @@ def evolve_mass_fraction(raw_isotope_mass_fraction, time_array):
     return time_evolved_isotope_mass_fraction
 
 
-def time_evolve_mass_fractions(
+def time_evolve_cumulative_decay(
     raw_isotope_mass_fraction, shell_masses, gamma_ray_lines, time_array
 ):
     """
@@ -231,35 +231,29 @@ def time_evolve_mass_fractions(
 
     """
 
-    isotope_mass_fraction_list = []
-    cumulative_decay_df_list = []
+    isotope_decay_df_list = []
     initial_isotope_mass_fraction = raw_isotope_mass_fraction
 
-    decay_times = np.diff(time_array)
+    dt = np.diff(time_array)
 
-    for time in decay_times:
-        decayed_isotope_mass_fraction = IsotopicMassFraction(
-            initial_isotope_mass_fraction
-        ).decay(time)
+    for time in dt:
 
         isotope_dict = create_isotope_dicts(
-            decayed_isotope_mass_fraction, shell_masses
+            initial_isotope_mass_fraction, shell_masses
         )
         inventories = create_inventories_dict(isotope_dict)
-        cumulative_decay_df = calculate_total_decays(inventories, time)
-        isotope_decay_df = create_isotope_decay_df(
-            cumulative_decay_df, gamma_ray_lines
-        )
-        isotope_mass_fraction_list.append(isotope_decay_df)
-        cumulative_decay_df_list.append(cumulative_decay_df)
-        initial_isotope_mass_fraction = decayed_isotope_mass_fraction
+        total_decays = calculate_total_decays(inventories, time)
+        isotope_df_time = create_isotope_decay_df(total_decays, gamma_ray_lines)
+        isotope_decay_df_list.append(isotope_df_time)
 
-    time_evolved_cumulative_decay = pd.concat(
-        cumulative_decay_df_list, keys=time_array, names=["time"]
-    )
+        decayed_isotope_mass_fraction = IsotopicMassFraction(
+            initial_isotope_mass_fraction
+        ).decay(time)
 
-    time_evolve_decay_df = pd.concat(
-        isotope_mass_fraction_list, keys=time_array, names=["time"]
+        initial_isotope_mass_fraction = decayed_isotope_mass_fraction
+
+    time_evolved_decay_df = pd.concat(
+        isotope_decay_df_list, keys=time_array, names=["time"]
     )
 
-    return time_evolve_decay_df, time_evolved_cumulative_decay
+    return time_evolved_decay_df

From 960f2b8915acac71a553824ca8bc762d3941c8f6 Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Wed, 17 Jul 2024 10:38:35 -0400
Subject: [PATCH 16/18] Added tests for cumulative decays

---
 .../energy_input/gamma_ray_packet_source.py   | 24 ++++++------
 tardis/energy_input/main_gamma_ray_loop.py    |  5 +--
 .../tests/test_gamma_ray_channel.py           | 39 +++++++++++++++++++
 3 files changed, 52 insertions(+), 16 deletions(-)

diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index 51caf32ba89..6a150c87a5a 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -437,18 +437,18 @@ def create_packets(self, decays_per_isotope, *args, **kwargs):
                 packet_energies_cmf[
                     packet_index - isotope_packet_count : packet_index
                 ] = initial_packet_energies_cmf
-                nus_rf[packet_index - isotope_packet_count : packet_index] = (
-                    initial_nus_rf
-                )
-                nus_cmf[packet_index - isotope_packet_count : packet_index] = (
-                    initial_nus_cmf
-                )
-                shells[packet_index - isotope_packet_count : packet_index] = (
-                    initial_shells
-                )
-                times[packet_index - isotope_packet_count : packet_index] = (
-                    initial_times
-                )
+                nus_rf[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_nus_rf
+                nus_cmf[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_nus_cmf
+                shells[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_shells
+                times[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_times
 
         return GXPacketCollection(
             locations,
diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py
index 6f8fad0bcac..002f50301b3 100644
--- a/tardis/energy_input/main_gamma_ray_loop.py
+++ b/tardis/energy_input/main_gamma_ray_loop.py
@@ -241,10 +241,7 @@ def run_gamma_ray_loop(
     logger.info(f"Total CMF energy is {total_cmf_energy}")
     logger.info(f"Total RF energy is {total_rf_energy}")
 
-    (
-        energy_out,
-        packets_array,
-    ) = gamma_packet_loop(
+    (energy_out, packets_array,) = gamma_packet_loop(
         packets,
         grey_opacity,
         photoabsorption_opacity,
diff --git a/tardis/energy_input/tests/test_gamma_ray_channel.py b/tardis/energy_input/tests/test_gamma_ray_channel.py
index ea8e16d034b..372e2816801 100644
--- a/tardis/energy_input/tests/test_gamma_ray_channel.py
+++ b/tardis/energy_input/tests/test_gamma_ray_channel.py
@@ -17,9 +17,11 @@
     create_inventories_dict,
     calculate_total_decays,
     create_isotope_decay_df,
+    time_evolve_cumulative_decay,
 )
 from tardis.energy_input.gamma_ray_transport import get_taus
 from tardis.energy_input.util import KEV2ERG
+from tardis.energy_input.main_gamma_ray_loop import get_effective_time_array
 
 
 @pytest.fixture(scope="module")
@@ -313,3 +315,40 @@ def test_total_energy_production(gamma_ray_test_composition, atomic_dataset):
     actual = ni56_energy + co56_energy
 
     npt.assert_allclose(actual, expected)
+
+
+def test_cumulative_decays(gamma_ray_test_composition, atomic_dataset):
+    """
+    Function to test the cumulative decay of isotopes.
+    Parameters
+    ----------
+    gamma_ray_simulation_state: Tardis simulation state
+    atomic_dataset: Tardis atomic-nuclear dataset
+    """
+
+    time_start = 0.1 * u.d
+    time_end = 100 * u.d
+    time_steps = 2
+    time_space = "linear"
+    time_delta = (time_end - time_start).value
+
+    gamma_ray_lines = atomic_dataset.decay_radiation_data
+    raw_isotopic_mass_fraction, cell_masses = gamma_ray_test_composition
+    isotope_dict = create_isotope_dicts(raw_isotopic_mass_fraction, cell_masses)
+    inventories_dict = create_inventories_dict(isotope_dict)
+    total_decays = calculate_total_decays(inventories_dict, time_delta)
+    isotope_decay_df = create_isotope_decay_df(total_decays, gamma_ray_lines)
+
+    times, effective_times = get_effective_time_array(
+        time_start.value, time_end.value, time_space, time_steps
+    )
+    # total decay energy in the entire time range
+    actual = isotope_decay_df["decay_energy_erg"].sum()
+
+    # time evolve the decay energy
+    evolve_decays_with_time = time_evolve_cumulative_decay(
+        raw_isotopic_mass_fraction, cell_masses, gamma_ray_lines, times
+    )
+    expected = evolve_decays_with_time["decay_energy_erg"].sum()
+
+    npt.assert_allclose(actual, expected)

From 6161f6a0cff4ec8545a40899020b8e56818d8e31 Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Mon, 22 Jul 2024 16:38:40 -0400
Subject: [PATCH 17/18] Added test for time evolution of decays as the
 composition changes

---
 .../energy_input/gamma_ray_packet_source.py   | 101 +++---------------
 .../tests/test_gamma_ray_channel.py           |  12 ++-
 2 files changed, 21 insertions(+), 92 deletions(-)

diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index 6a150c87a5a..59eccb6493b 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -437,18 +437,18 @@ def create_packets(self, decays_per_isotope, *args, **kwargs):
                 packet_energies_cmf[
                     packet_index - isotope_packet_count : packet_index
                 ] = initial_packet_energies_cmf
-                nus_rf[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_nus_rf
-                nus_cmf[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_nus_cmf
-                shells[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_shells
-                times[
-                    packet_index - isotope_packet_count : packet_index
-                ] = initial_times
+                nus_rf[packet_index - isotope_packet_count : packet_index] = (
+                    initial_nus_rf
+                )
+                nus_cmf[packet_index - isotope_packet_count : packet_index] = (
+                    initial_nus_cmf
+                )
+                shells[packet_index - isotope_packet_count : packet_index] = (
+                    initial_shells
+                )
+                times[packet_index - isotope_packet_count : packet_index] = (
+                    initial_times
+                )
 
         return GXPacketCollection(
             locations,
@@ -708,73 +708,6 @@ def create_packet_times_uniform_energy(
                 )
         return decay_times
 
-    def sample_decay_times(self, isotopes, taus, times, number_of_packets):
-        """
-        Sample decay times for a given number of packets
-
-        Parameters
-        ----------
-        isotopes : array
-            Array of isotope names as strings
-        taus : dict
-            Dictionary of isotope mean lifetimes in seconds
-        times : array
-            Array of time steps in seconds
-        number_of_packets : int
-            Number of packets
-
-        Returns
-        -------
-        decay_times : array
-            Array of decay times for each packet
-
-        """
-
-        decay_times = np.zeros(number_of_packets)
-
-        for i, isotope in enumerate(isotopes):
-            # decay time of the packet cannot be less than the minimum time in the simulation
-            decay_time_min = times[0]
-            # decay time of the packet cannot be greater than the maximum time in the simulation
-            decay_time_max = times[-1]
-
-            # rejection sampling on decay times
-            while (decay_times[i] <= decay_time_min) or (
-                decay_times[i] >= decay_time_max
-            ):
-                decay_times[i] = -taus[isotope] * np.log(np.random.random())
-
-        return decay_times
-
-    def sample_decay_times_uniformly(self, time, number_of_packets, seed):
-        """
-        Sample decay times uniformly
-
-        Parameters
-        ----------
-        time : array
-            Array of time steps in seconds
-        number_of_packets : int
-            Number of packets
-
-        Returns
-        -------
-        decay_times : array
-            Array of decay times for each packet
-
-        """
-
-        decay_times = np.zeros(number_of_packets)
-
-        np.random.seed(seed)
-        decay_times[0] = self.times[0]
-        decay_times[-1] = self.times[-1]
-
-        for j in range(1, number_of_packets - 1):
-            decay_times[j] = np.random.uniform(time[j - 1], time[j])
-
-        return decay_times
-
     def create_packets(
         self, decays_per_isotope, number_of_packets, seed, *args, **kwargs
     ):
@@ -833,16 +766,6 @@ def create_packets(
             sampled_packets_df.index.get_level_values("time") * 86400.0
         )
 
-        # Get the indices of the time steps from the sampled times
-
-        # decay_times = self.sample_decay_times_uniformly(
-        #    sampled_times, number_of_packets, seed
-        # )
-
-        # get the time of the middle of the step for each packet
-        # decay_time_indices = np.searchsorted(self.effective_times, decay_times)
-        # 3D locations
-
         decay_time_indices = []
         for i in range(number_of_packets):
             decay_time_indices.append(
diff --git a/tardis/energy_input/tests/test_gamma_ray_channel.py b/tardis/energy_input/tests/test_gamma_ray_channel.py
index 372e2816801..f25028e545f 100644
--- a/tardis/energy_input/tests/test_gamma_ray_channel.py
+++ b/tardis/energy_input/tests/test_gamma_ray_channel.py
@@ -319,7 +319,10 @@ def test_total_energy_production(gamma_ray_test_composition, atomic_dataset):
 
 def test_cumulative_decays(gamma_ray_test_composition, atomic_dataset):
     """
-    Function to test the cumulative decay of isotopes.
+    Function to test that the total energy calculated from summing all the decays
+    from the entire time range of simulation is the same as decay energy from individual
+    time steps considering that after each time step the composition (mass fractions) changes.
+    Tested for Ni56, Cr48, Fe52.
     Parameters
     ----------
     gamma_ray_simulation_state: Tardis simulation state
@@ -328,7 +331,7 @@ def test_cumulative_decays(gamma_ray_test_composition, atomic_dataset):
 
     time_start = 0.1 * u.d
     time_end = 100 * u.d
-    time_steps = 2
+    time_steps = 3
     time_space = "linear"
     time_delta = (time_end - time_start).value
 
@@ -351,4 +354,7 @@ def test_cumulative_decays(gamma_ray_test_composition, atomic_dataset):
     )
     expected = evolve_decays_with_time["decay_energy_erg"].sum()
 
-    npt.assert_allclose(actual, expected)
+    # This rtol is set since the decay energy is calculated with Fe52 (which has Mn-52m as a daughter)
+    # The data is not available for Mn-52m in the decay_radiation_data
+    # If we use any other isotope without a metastable state, the total decay energy matches exactly.
+    npt.assert_allclose(actual, expected, rtol=1e-4)

From 7b123f848df55192f3452536d75d02237ee03c5b Mon Sep 17 00:00:00 2001
From: Knights-Templars <anirbaniamdutta@gmail.com>
Date: Mon, 22 Jul 2024 16:52:23 -0400
Subject: [PATCH 18/18] Ran black

---
 .../energy_input/gamma_ray_packet_source.py   | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py
index 59eccb6493b..0374cdf7788 100644
--- a/tardis/energy_input/gamma_ray_packet_source.py
+++ b/tardis/energy_input/gamma_ray_packet_source.py
@@ -437,18 +437,18 @@ def create_packets(self, decays_per_isotope, *args, **kwargs):
                 packet_energies_cmf[
                     packet_index - isotope_packet_count : packet_index
                 ] = initial_packet_energies_cmf
-                nus_rf[packet_index - isotope_packet_count : packet_index] = (
-                    initial_nus_rf
-                )
-                nus_cmf[packet_index - isotope_packet_count : packet_index] = (
-                    initial_nus_cmf
-                )
-                shells[packet_index - isotope_packet_count : packet_index] = (
-                    initial_shells
-                )
-                times[packet_index - isotope_packet_count : packet_index] = (
-                    initial_times
-                )
+                nus_rf[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_nus_rf
+                nus_cmf[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_nus_cmf
+                shells[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_shells
+                times[
+                    packet_index - isotope_packet_count : packet_index
+                ] = initial_times
 
         return GXPacketCollection(
             locations,