Decay energy chain

tardis/energy_input/gamma_ray_transport.py

@@ -178,7 +178,6 @@ def initialize_packets(
 
     packet_index = 0
     for k, shell in enumerate(decays_per_shell):
-
         initial_radii = initial_packet_radius(
             shell, inner_velocities[k], outer_velocities[k]
         )
@@ -261,7 +260,6 @@ def initialize_packets(
 
 
 def calculate_shell_masses(model):
-
     """Function to calculate shell masses
     Parameters
     ----------
@@ -274,11 +272,206 @@ def calculate_shell_masses(model):
 
     """
 
-    ejecta_density = model.density.to("g/cm^3").value
-    ejecta_volume = model.volume.to("cm^3").value
-    shell_masses = ejecta_volume * ejecta_density
+    ejecta_density = model.density.to("g/cm^3")
+    ejecta_volume = model.volume.to("cm^3")
+    return (ejecta_volume * ejecta_density).to(u.g)
+
+
+def calculate_total_decays(inventories, time_delta):
+    """Function to create inventories of isotope
+    Parameters
+    ----------
+    model : tardis.Radial1DModel
+        The tardis model to calculate gamma ray propagation through
+
+    time_end : float
+        End time of simulation in days
+    Returns
+    -------
+        Total decay list : List
+            list of total decays for x g of isotope for time 't'
+
+    """
+
+    time_delta = u.Quantity(time_delta, u.s)
+
+    total_decays_list = []
+    for inv in inventories:
+        total_decays = inv.cumulative_decays(time_delta.value)
+        total_decays_list.append(total_decays)
+
+    return total_decays_list
+
+
+def create_isotope_dicts(raw_isotope_abundance, shell_masses):
+    """
+    Function to create a dictionary of isotopes for each shell with their masses.
+    Parameters
+    ----------
+    raw_isotope_abundance : pd.DataFrame
+        isotope abundance in mass fractions.
+
+    """
+    isotope_dicts = {}
+    for i in range(len(raw_isotope_abundance.columns)):
+        isotope_dicts[i] = {}
+        for (
+            atomic_number,
+            mass_number,
+        ), abundances in raw_isotope_abundance.iterrows():
+            isotope_dicts[i][atomic_number, mass_number] = {}
+            nuclear_symbol = f"{rd.utils.Z_to_elem(atomic_number)}{mass_number}"
+            isotope_dicts[i][atomic_number, mass_number][nuclear_symbol] = (
+                abundances[i] * shell_masses[i].to(u.g).value
+            )
+
+    return isotope_dicts
+
+
+def create_inventories_dict(isotope_dict):
+    inv = {}
+    for shell, isotopes in isotope_dict.items():
+        inv[shell] = {}
+        for isotope, name in isotopes.items():
+            inv[shell][isotope] = rd.Inventory(name, "g")
+
+    return inv
+
+
+def calculate_total_decays(inventory_dict, time_delta):
+    time_delta = u.Quantity(time_delta, u.s)
+    total_decays = {}
+    for shell, isotopes in inventory_dict.items():
+        total_decays[shell] = {}
+        for isotope, name in isotopes.items():
+            total_decays[shell][isotope] = name.cumulative_decays(
+                time_delta.value
+            )
+
+    return total_decays
+
+
+def calculate_average_energies(raw_isotope_abundance, gamma_ray_lines):
+    """
+    Function to calculate average energies of positrons and gamma rays
+    from a list of gamma ray lines from nndc.
+    Parameters
+    ----------
+    raw_isotope_abundance : pd.DataFrame
+        isotope abundance in mass fractions
+    gamma_ray_lines : pd.DataFrame
+        decay data
+
+    Returns
+    -------
+    average_energies_list : List
+        list of gamma ray energies
+    average_positron_energies_list : List
+        list of positron energies
+    gamma_ray_line_array_list : List
+        list of gamma ray lines
+
+    """
+
+    all_isotope_names = get_all_isotopes(raw_isotope_abundance)
+    all_isotope_names.sort()
+
+    gamma_ray_line_array_list = []
+    average_energies_list = []
+    average_positron_energies_list = []
+
+    for i, isotope in enumerate(all_isotope_names):
+        energy, intensity = setup_input_energy(
+            gamma_ray_lines[gamma_ray_lines.index == isotope.replace("-", "")],
+            "g",
+        )
+        average_energies_list.append(np.sum(energy * intensity))  # keV
+        gamma_ray_line_array_list.append(np.stack([energy, intensity]))
+
+        positron_energy, positron_intensity = setup_input_energy(
+            gamma_ray_lines[gamma_ray_lines.index == isotope.replace("-", "")],
+            "bp",
+        )
+        average_positron_energies_list.append(
+            np.sum(positron_energy * positron_intensity)
+        )
+
+    return (
+        average_energies_list,
+        average_positron_energies_list,
+        gamma_ray_line_array_list,
+    )
+
+
+def decay_chain_energies(
+    raw_isotope_abundance,
+    average_energies_list,
+    average_positron_energies_list,
+    gamma_ray_line_array_list,
+    total_decays,
+):
+    all_isotope_names = get_all_isotopes(raw_isotope_abundance)
+    all_isotope_names.sort()
+
+    gamma_ray_line_arrays = {}
+    average_energies = {}
+    average_positron_energies = {}
+
+    for iso, lines in zip(all_isotope_names, gamma_ray_line_array_list):
+        gamma_ray_line_arrays[iso] = lines
+
+    for iso, energy, positron_energy in zip(
+        all_isotope_names, average_energies_list, average_positron_energies_list
+    ):
+        average_energies[iso] = energy
+        average_positron_energies[iso] = positron_energy
+
+    decay_energy = {}
+    for shell, isotopes in total_decays.items():
+        decay_energy[shell] = {}
+        for name, isotope in isotopes.items():
+            decay_energy[shell][name] = {}
+            for iso, dps in isotope.items():
+                # print(iso)
+                decay_energy[shell][name][iso] = dps * average_energies[iso]
+
+    return decay_energy
+
+
+def calculate_energy_per_mass(
+    decay_energy, raw_isotope_abundance, shell_masses
+):
+    energy_dict = {}
+    for shell, isotopes in decay_energy.items():
+        energy_dict[shell] = {}
+        for name, isotope in isotopes.items():
+            energy_dict[shell][name] = sum(isotope.values())
+
+    energy_list = []
+    for shell, isotopes in energy_dict.items():
+        for isotope, energy in isotopes.items():
+            energy_list.append(
+                {
+                    "shell": shell,
+                    "atomic_number": isotope[0],
+                    "mass_number": isotope[1],
+                    "value": energy,
+                }
+            )
+
+    df = pd.DataFrame(energy_list)
+    energy_df = pd.pivot_table(
+        df,
+        values="value",
+        index=["atomic_number", "mass_number"],
+        columns="shell",
+    )
+
+    energy_per_mass = energy_df.divide(
+        (raw_isotope_abundance * shell_masses).to_numpy(), axis=0
+    )
 
-    return shell_masses
+    return energy_per_mass, energy_df
 
 
 def main_gamma_ray_loop(
@@ -368,6 +561,7 @@ def main_gamma_ray_loop(
     )
 
     shell_masses = calculate_shell_masses(model)
+    inventories = raw_isotope_abundance.to_inventories(shell_masses)
 
     time_start = time_explosion
     time_end *= u.d.to(u.s)
@@ -427,7 +621,6 @@ def main_gamma_ray_loop(
     number_of_isotopes = plasma.isotope_number_density * ejecta_volume
     total_number_isotopes = number_of_isotopes.sum(axis=1)
 
-    inventories = raw_isotope_abundance.to_inventories()
     all_isotope_names = get_all_isotopes(raw_isotope_abundance)
     all_isotope_names.sort()