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]
         )
@@ -274,11 +273,199 @@ 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
 
-    return shell_masses
+    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):
+    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 energy_per_mass, energy_df
 
 
 def main_gamma_ray_loop(
@@ -368,6 +555,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 +615,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()
 

tardis/energy_input/tests/test_gamma_ray_transport.py

@@ -0,0 +1,155 @@
+import os
+import pytest
+import numpy as np
+import numpy.testing as npt
+from pathlib import Path
+import radioactivedecay as rd
+
+import tardis
+from tardis.model import SimulationState
+from tardis.io.configuration import config_reader
+from tardis.io.util import yaml_load_file, YAMLLoader
+from tardis.energy_input.gamma_ray_transport import (
+    calculate_shell_masses,
+    create_isotope_dicts,
+    get_all_isotopes,
+    create_inventories_dict,
+    calculate_total_decays,
+)
+import astropy.units as u
+import astropy.constants as c
+
+# DATA_PATH = os.path.join(
+#    tardis.__path__[0], "io", "configuration", "tests", "data"
+# )
+NI56_NUCLIDE = rd.Nuclide("Ni56")
+
+
+@pytest.fixture(scope="module")
+def gamma_ray_config(example_configuration_dir: Path):
+    """
+    Parameters
+    ----------
+    example_configuration_dir: Path to the configuration directory.
+
+    Returns
+    -------
+    Tardis configuration
+    """
+    yml_path = (
+        example_configuration_dir
+        / "tardis_configv1_density_exponential_nebular.yml"
+    )
+
+    return config_reader.Configuration.from_yaml(yml_path)
+
+
+@pytest.fixture(scope="module")
+def simulation_setup(gamma_ray_config):
+    """
+    Parameters
+    ----------
+    gamma_ray_config:
+
+    Returns
+    -------
+
+    """
+    gamma_ray_config.model.structure.velocity.start = 1.0 * u.km / u.s
+    gamma_ray_config.model.structure.density.rho_0 = 5.0e2 * u.g / (u.cm**3)
+    gamma_ray_config.supernova.time_explosion = 1.0 * (u.d)
+    model = SimulationState.from_config(gamma_ray_config)
+    return model
+
+
+def test_calculate_shell_masses(simulation_setup):
+    """
+    Function to test calculation of shell masses.
+    Parameters
+    ----------
+    simulation_setup: A simulation setup which returns a model.
+    """
+    model = simulation_setup
+    volume = model.volume.to("cm^3")
+    density = model.density.to("g/cm^3")
+    desired = (volume * density).to(u.g).value
+
+    shell_masses = calculate_shell_masses(model).value
+    npt.assert_almost_equal(shell_masses, desired)
+
+
+@pytest.mark.parametrize("nuclide_name", ["Ni-56"])
+def test_activity(simulation_setup, nuclide_name):
+    """
+    Function to test the decay of 56Ni in radioactivedecay with an analytical solution.
+    Parameters
+    ----------
+    simulation_setup: A simulation setup which returns a model.
+    nuclide_name: Name of the nuclide.
+    """
+    nuclide = rd.Nuclide(nuclide_name)
+    model = simulation_setup
+    t_half = nuclide.half_life() * u.s
+    decay_constant = np.log(2) / t_half
+    time_delta = 1.0 * u.s
+    shell_masses = calculate_shell_masses(model)
+    raw_isotope_abundance = model.raw_isotope_abundance
+    raw_isotope_abundance_mass = raw_isotope_abundance.apply(
+        lambda x: x * shell_masses, axis=1
+    )
+    mass = raw_isotope_abundance_mass.loc[nuclide.Z, nuclide.A][0]
+    iso_dict = create_isotope_dicts(raw_isotope_abundance, shell_masses)
+    inv_dict = create_inventories_dict(iso_dict)
+
+    total_decays = calculate_total_decays(inv_dict, time_delta)
+    actual = total_decays[0][nuclide.Z, nuclide.A][nuclide_name]
+
+    isotopic_mass = nuclide.atomic_mass * (u.g)
+    number_of_moles = mass * (u.g) / isotopic_mass
+    number_of_atoms = (number_of_moles * c.N_A).value
+    N1 = number_of_atoms * np.exp(-decay_constant * time_delta)
+    expected = number_of_atoms - N1.value
+
+    npt.assert_allclose(actual, expected, rtol=1e-7)
+
+
+@pytest.mark.xfail(reason="To be implemented")
+def test_activity_chain(simulation_setup):
+    model = simulation_setup
+    t_half_Ni = NI_INV.half_lives("s")["Ni-56"]
+    t_half_Co = 77.236 * u.d.to(u.s)
+    decay_constant_Ni = np.log(2) / t_half_Ni
+    decay_constant_Co = np.log(2) / t_half_Co
+    time_delta = 80.0 * u.d.to(u.s)
+    shell_masses = calculate_shell_masses(model)
+    raw_isotope_abundance = model.raw_isotope_abundance
+    raw_isotope_abundance_mass = raw_isotope_abundance.apply(
+        lambda x: x * shell_masses, axis=1
+    )
+    ni_mass = raw_isotope_abundance_mass.loc[28, 56][0]
+    iso_dict = create_isotope_dicts(raw_isotope_abundance, shell_masses)
+    inv_dict = create_inventories_dict(iso_dict)
+    total_decays = calculate_total_decays(inv_dict, time_delta)
+    actual_Ni = total_decays[0][28, 56]["Ni-56"]
+    actual_Co = total_decays[0][28, 56]["Co-56"]
+
+    isotopic_mass_Ni = NI_INV._get_atomic_mass("Ni-56") * (u.g)
+    number_of_moles = ni_mass * (u.g) / isotopic_mass_Ni
+    number_of_atoms = number_of_moles * c.N_A
+    N1 = number_of_atoms.value * np.exp(-decay_constant_Ni * time_delta)
+    N2 = (
+        number_of_atoms.value
+        * decay_constant_Ni
+        / (decay_constant_Co - decay_constant_Ni)
+        * (
+            np.exp(-decay_constant_Ni * time_delta)
+            - np.exp(-decay_constant_Co * time_delta)
+        )
+    )
+    expected_Ni = number_of_atoms.value * (
+        1 - np.exp(-decay_constant_Ni * time_delta)
+    )
+    expected_Co = number_of_atoms.value - N1 - N2
+
+    npt.assert_almost_equal(actual_Ni, expected_Ni)
+    npt.assert_almost_equal(actual_Co, expected_Co)