Changed the positron fraction and updated with a function to extract …

…packet info

tardis/energy_input/gamma_ray_packet_source.py

@@ -740,12 +740,22 @@ def create_packets(
         nus_cmf = np.zeros(number_of_packets)
         statuses = np.ones(number_of_packets, dtype=np.int64) * 3
 
-        # sample packets from the gamma-ray lines
+        # sample packets from the gamma-ray lines only (include X-rays!)
         sampled_packets_df_gamma = decays_per_isotope[
             decays_per_isotope["radiation"] == "g"
         ]
-        # sample packets from dataframe, returning a dataframe where each row is
-        # a sampled packet
+
+        # sample packets from the gamma-ray lines at time = t0 : which is the start of the gamma-ray simulation
+        #sampled_packets_df_t0 = sampled_packets_df_gamma[sampled_packets_df_gamma.index.get_level_values("time") == self.times[0] * (u.s).to(u.d)]
+        # sample same number of packets as the gamma-ray lines at time = 0
+        # sampled_packets_df_gamma_t0 = sampled_packets_df_t0.sample(
+        #     n=number_of_packets,
+        #     weights="decay_energy_erg",
+        #     replace=True,
+        #     random_state=np.random.RandomState(self.base_seed),
+        # )
+
+        # sample packets from the time evolving dataframe
         sampled_packets_df = sampled_packets_df_gamma.sample(
             n=number_of_packets,
             weights="decay_energy_erg",
@@ -765,17 +775,21 @@ def create_packets(
         sampled_packets_df["inner_velocity"] = self.inner_velocities[shells]
         sampled_packets_df["outer_velocity"] = self.outer_velocities[shells]
 
-        # sample radii at time = 0
+        # 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") * (
             u.d
         ).to(u.s)
         # TODO: Rewrite this without for loop. This is expensive
+        # get the time step index of the packets
         decay_time_indices = []
         for i in range(number_of_packets):
             decay_time_indices.append(get_index(sampled_times[i], self.times))
 
+        # 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.
         locations = (
             initial_radii.values
             * self.effective_times[decay_time_indices]
@@ -843,21 +857,24 @@ def calculate_positron_fraction(self, isotopes, number_of_packets):
 
         # 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
+           self.isotope_decay_df.index.get_level_values("shell_number") == 0
         ]
 
+        gamma_decay_df = shell_number_0[shell_number_0["radiation"] == "g"]
+
         positrons_decay_df = shell_number_0[shell_number_0["radiation"] == "bp"]
+        # Find the total energy released from positrons per isotope from the dataframe
         positron_energy_per_isotope = positrons_decay_df.groupby("isotope")[
             "energy_per_channel_keV"
         ].sum()
-        # Find the total energy released per isotope from the dataframe
-        total_energy_per_isotope = shell_number_0.groupby("isotope")[
+        # Find the total energy released from gamma-ray per isotope from the dataframe
+        gamma_energy_per_isotope = gamma_decay_df.groupby("isotope")[
             "energy_per_channel_keV"
         ].sum()
         # TODO: Possibly move this for loop
         for i, isotope in enumerate(isotopes):
             isotope_positron_fraction[i] = (
                 positron_energy_per_isotope[isotope]
-                / total_energy_per_isotope[isotope]
+                / gamma_energy_per_isotope[isotope]
             )
         return isotope_positron_fraction

tardis/energy_input/main_gamma_ray_loop.py

@@ -12,6 +12,7 @@
     iron_group_fraction_per_shell,
 )
 from tardis.energy_input.GXPacket import GXPacket
+from tardis.energy_input.util import get_index
 from tardis.energy_input.util import make_isotope_string_tardis_like
 
 logger = logging.getLogger(__name__)
@@ -207,8 +208,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))
     energy_deposited = np.zeros((number_of_shells, time_steps))
@@ -222,8 +221,8 @@ def run_gamma_ray_loop(
     total_rf_energy = 0
 
     for p in packets:
-        total_cmf_energy += p.energy_cmf
-        total_rf_energy += p.energy_rf
+       total_cmf_energy += p.energy_cmf
+       total_rf_energy += p.energy_rf
 
     logger.info(f"Total CMF energy is {total_cmf_energy}")
     logger.info(f"Total RF energy is {total_rf_energy}")
@@ -279,8 +278,62 @@ def run_gamma_ray_loop(
     positron_energy = pd.DataFrame(
         data=energy_deposited_positron, columns=times[:-1]
     )
-    print(positron_energy.sum().sum())
 
     total_deposited_energy = (positron_energy + deposited_energy) / dt_array
 
-    return escape_energy, packets_df_escaped, total_deposited_energy
+    return escape_energy, packets_df_escaped, deposited_energy, total_deposited_energy,
+
+
+
+def get_packet_properties(number_of_shells, times, time_steps, packets):
+
+    """
+    Function to get the properties of the packets.
+
+    Parameters
+    ----------
+    packets : list
+        List of packets.
+    
+    Returns
+    -------
+    packets_nu_cmf_array : np.ndarray
+        Array of packets in cmf.
+    packets_nu_rf_array : np.ndarray
+        Array of packets in rf.
+    packets_energy_cmf_array : np.ndarray
+        Array of packets energy in cmf.
+    packets_energy_rf_array : np.ndarray
+        Array of packets energy in rf.
+    packets_positron_energy_array : np.ndarray
+        Array of packets positron energy.
+    """
+
+    #collect the properties of the packets
+    shell_number = []
+    time_current = []
+
+     # Bin the frequency of the packets in shell and time
+
+    packets_nu_cmf_array = np.zeros((number_of_shells, time_steps))
+    packets_nu_rf_array = np.zeros((number_of_shells, time_steps))
+    packets_energy_cmf_array = np.zeros((number_of_shells, time_steps))
+    packets_energy_rf_array = np.zeros((number_of_shells, time_steps))
+    packets_positron_energy_array = np.zeros((number_of_shells, time_steps))
+
+
+    for p in packets:
+        time_index = get_index(p.time_current, times)
+        shell_number.append(p.shell)
+        time_current.append(p.time_current)
+        packets_nu_cmf_array[p.shell, time_index] += p.nu_cmf
+        packets_nu_rf_array[p.shell, time_index] += p.nu_rf
+        packets_energy_cmf_array[p.shell, time_index] += p.energy_cmf
+        packets_energy_rf_array[p.shell, time_index] += p.energy_rf
+        packets_positron_energy_array[p.shell, time_index] += p.positron_energy
+
+
+    return packets_nu_cmf_array, packets_nu_rf_array, packets_energy_cmf_array, packets_energy_rf_array, packets_positron_energy_array
+
+
+