From 7b7576190c42949331eafe7443d6b9caef842dc1 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 18 Mar 2024 14:53:30 -0400 Subject: [PATCH] Positron deposition and packet tracker arrays --- tardis/energy_input/GXPacket.py | 7 -- tardis/energy_input/gamma_packet_loop.py | 5 -- .../energy_input/gamma_ray_packet_source.py | 28 +++++-- tardis/energy_input/gamma_ray_transport.py | 17 ++++- tardis/energy_input/main_gamma_ray_loop.py | 74 ++++++++----------- 5 files changed, 67 insertions(+), 64 deletions(-) diff --git a/tardis/energy_input/GXPacket.py b/tardis/energy_input/GXPacket.py index 661b593a2bd..745aa850184 100644 --- a/tardis/energy_input/GXPacket.py +++ b/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 diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py index cd4307ffba9..8f76d1f2c07 100644 --- a/tardis/energy_input/gamma_packet_loop.py +++ b/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, ) diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py index d203460faea..71d274091b7 100644 --- a/tardis/energy_input/gamma_ray_packet_source.py +++ b/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 diff --git a/tardis/energy_input/gamma_ray_transport.py b/tardis/energy_input/gamma_ray_transport.py index ee847762c9a..45d662e5d16 100644 --- a/tardis/energy_input/gamma_ray_transport.py +++ b/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] } diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py index 459a5c958ba..3d8bd9badf0 100644 --- a/tardis/energy_input/main_gamma_ray_loop.py +++ b/tardis/energy_input/main_gamma_ray_loop.py @@ -94,7 +94,9 @@ 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 @@ -102,7 +104,9 @@ def run_gamma_ray_loop( 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, )