Positron deposition and packet tracker arrays

tardis/energy_input/GXPacket.py

@@ -104,13 +104,6 @@ def __init__(
         self.number_of_packets = len(self.energy_rf)
         self.tau = -np.log(np.random.random(self.number_of_packets))
 
-    def get_energy_plot_df(self):
-        energy_plot_df_rows = np.zeros((self.number_of_packets, 8))
-        pass
-
-    def get_energy_df(self):
-        pass
-
     def get_positron_energy_df(self):
         energy_plot_positron_rows = np.zeros((self.number_of_packets, 4))
         pass

tardis/energy_input/gamma_packet_loop.py

@@ -50,7 +50,6 @@ def gamma_packet_loop(
     energy_df_rows,
     energy_plot_df_rows,
     energy_out,
-    packets_out,
     packets_info_array,
 ):
     """Propagates packets through the simulation
@@ -289,9 +288,6 @@ def gamma_packet_loop(
                     energy_out[bin_index, time_index] += rest_energy / (
                         bin_width * dt
                     )
-                    packets_out[bin_index] = np.array(
-                        [bin_index, i, rest_energy, packet.Z, packet.A]
-                    )
                     packet.status = GXPacketStatus.ESCAPED
                     escaped_packets += 1
                     if scattered:
@@ -322,7 +318,6 @@ def gamma_packet_loop(
         energy_plot_df_rows,
         energy_out,
         deposition_estimator,
-        packets_out,
         bin_width,
         packets_info_array,
     )

tardis/energy_input/gamma_ray_packet_source.py

@@ -1,5 +1,4 @@
 import numpy as np
-from numba import njit
 
 from tardis.energy_input.energy_source import (
     positronium_continuum,
@@ -14,7 +13,6 @@
     get_index,
     get_random_unit_vector,
 )
-from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel
 from tardis.montecarlo.packet_source import BasePacketSource
 
 
@@ -49,6 +47,7 @@ def __init__(
         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):
@@ -256,9 +255,6 @@ def create_packets(self, decays_per_isotope, *args, **kwargs):
         number_of_packets = decays_per_isotope.sum().sum()
         decays_per_shell = decays_per_isotope.T.sum().values
 
-        energy_plot_df_rows = np.zeros((number_of_packets, 8))
-        energy_plot_positron_rows = np.zeros((number_of_packets, 4))
-
         locations = np.zeros((3, number_of_packets))
         directions = np.zeros((3, number_of_packets))
         packet_energies_rf = np.zeros(number_of_packets)
@@ -271,6 +267,8 @@ def create_packets(self, decays_per_isotope, *args, **kwargs):
 
         positronium_energy, positronium_intensity = positronium_continuum()
 
+        self.energy_plot_positron_rows = np.zeros((number_of_packets, 4))
+
         packet_index = 0
         # go through each shell
         for shell_number, pkts in enumerate(decays_per_shell):
@@ -383,8 +381,28 @@ def create_packets(self, decays_per_isotope, *args, **kwargs):
                         initial_packet_energies_cmf[i] / doppler_factor
                     )
                     initial_nus_rf[i] = initial_nus_cmf[i] / doppler_factor
+
+                    self.energy_plot_positron_rows[i] = np.array(
+                        [
+                            packet_index,
+                            isotope_positron_fraction
+                            * self.packet_energy
+                            * 1000,
+                            # * inv_volume_time[packet.shell, decay_time_index],
+                            initial_radii[i],
+                            initial_times[i],
+                        ]
+                    )
+
                     packet_index += 1
 
+                # deposit positron energy
+                for time in initial_time_indexes:
+                    self.energy_df_rows[shell_number, time] += (
+                        isotope_positron_fraction * self.packet_energy * 1000
+                    )
+
+                # collect packet properties
                 locations[
                     :, packet_index - isotope_packet_count : packet_index
                 ] = initial_locations

tardis/energy_input/gamma_ray_transport.py

@@ -58,7 +58,10 @@ def get_chain_decay_power_per_ejectamass(
 
         # this is the total decay energy from gamma-rays and positrons for the end of chain isotope
         # endecay = get_decaypath_lastnucdecayenergy(decaypathindex)
-        endecay = average_energies[end_isotope] + average_positron_energies[end_isotope]
+        endecay = (
+            average_energies[end_isotope]
+            + average_positron_energies[end_isotope]
+        )
 
         print("Decay energy, abundance, tau")
         print(endecay)
@@ -106,7 +109,9 @@ def calculate_total_decays(inventories, time_delta):
     for shell, isotopes in inventories.items():
         total_decays[shell] = {}
         for isotope, inventory in isotopes.items():
-            total_decays[shell][isotope] = inventory.cumulative_decays(time_delta.value)
+            total_decays[shell][isotope] = inventory.cumulative_decays(
+                time_delta.value
+            )
     return total_decays
 
 
@@ -397,7 +402,9 @@ def distribute_packets(decay_energy, energy_per_packet):
         for name, isotope in isotopes.items():
             packets_per_isotope[shell][name] = {}
             for line, energy in isotope.items():
-                packets_per_isotope[shell][name][line] = int(energy / energy_per_packet)
+                packets_per_isotope[shell][name][line] = int(
+                    energy / energy_per_packet
+                )
 
     packets_per_isotope_list = []
     for shell, parent_isotope in packets_per_isotope.items():
@@ -426,7 +433,9 @@ def packets_per_isotope(fractional_decay_energy, decayed_packet_count_dict):
     packets_per_isotope = {
         shell: {
             parent_isotope: {
-                isotopes: fractional_decay_energy[shell][parent_isotope][isotopes]
+                isotopes: fractional_decay_energy[shell][parent_isotope][
+                    isotopes
+                ]
                 * decayed_packet_count_dict[shell][parent_isotope]
                 for isotopes in fractional_decay_energy[shell][parent_isotope]
             }

tardis/energy_input/main_gamma_ray_loop.py

@@ -94,15 +94,19 @@ def run_gamma_ray_loop(
     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
+    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(raw_isotope_abundance, shell_masses)
     inventories_dict = create_inventories_dict(isotope_dict)
-    total_decays = calculate_total_decays(inventories_dict, time_end - time_start)
+    total_decays = calculate_total_decays(
+        inventories_dict, time_end - time_start
+    )
 
     (
         average_energies,
@@ -126,17 +130,6 @@ def run_gamma_ray_loop(
         total_isotope_number, axis=0
     )
 
-    # decayed_packet_count_dict = decayed_packet_count.to_dict()
-    # fractional_decay_energy_dict = fractional_decay_energy(decayed_energy)
-    # packets_per_isotope_df = (
-    #    packets_per_isotope(
-    #        fractional_decay_energy_dict, decayed_packet_count_dict
-    #    )
-    #    .round()
-    #    .fillna(0)
-    #    .astype(int)
-    # )
-
     total_energy = energy_df.sum().sum()
     energy_per_packet = total_energy / num_decays
     packets_per_isotope_df = (
@@ -176,7 +169,13 @@ def run_gamma_ray_loop(
 
     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 = 0.0
+    total_rf_energy = 0.0
     for i in range(number_of_packets):
         packet = GXPacket(
             packet_collection.location[:, i],
@@ -190,28 +189,33 @@ def run_gamma_ray_loop(
             packet_collection.time_current[i],
         )
         packets.append(packet)
-
-    total_cmf_energy = 0.0
-    total_rf_energy = 0.0
-
-    for p in packets:
-        total_cmf_energy += p.energy_cmf
-        total_rf_energy += p.energy_rf
+        energy_plot_df_rows[i] = np.array(
+            [
+                i,
+                packet.energy_rf,
+                packet.get_location_r(),
+                packet.time_current,
+                int(packet.status),
+                0,
+                0,
+                0,
+            ]
+        )
+        total_cmf_energy += packet.energy_cmf
+        total_rf_energy += packet.energy_rf
 
     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_out = np.zeros((len(energy_bins - 1), 5))
-    packets_info_array = np.zeros((num_decays, 10))
+    packets_info_array = np.zeros((int(num_decays), 8))
 
     (
         energy_df_rows,
         energy_plot_df_rows,
         energy_out,
         deposition_estimator,
-        packets_out,
         bin_width,
         packets_array,
     ) = gamma_packet_loop(
@@ -232,7 +236,6 @@ def run_gamma_ray_loop(
         energy_df_rows,
         energy_plot_df_rows,
         energy_out,
-        packets_out,
         packets_info_array,
     )
 
@@ -251,7 +254,7 @@ def run_gamma_ray_loop(
     )
 
     energy_plot_positrons = pd.DataFrame(
-        data=energy_plot_positron_rows,
+        data=packet_source.energy_plot_positron_rows,
         columns=[
             "packet_index",
             "energy_input",
@@ -260,28 +263,14 @@ def run_gamma_ray_loop(
         ],
     )
 
-    packets_out_df = pd.DataFrame(
-        data=packets_array,
-        columns=[
-            "number",
-            "status",
-            "Z",
-            "A",
-            "nu_cmf",
-            "nu_rf",
-            "energy_cmf",
-            "lum_rf",
-            "energy_rf",
-            "shell",
-        ],
-    )
-
     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)
+    escape_energy = pd.DataFrame(
+        data=energy_out, columns=times[:-1], index=energy_bins
+    )
 
     return (
         energy_df,
@@ -290,5 +279,4 @@ def run_gamma_ray_loop(
         decayed_packet_count,
         energy_plot_positrons,
         energy_estimated_deposition,
-        packets_out_df,
     )