From 38ee6bd35a6a839e998f550f5a59a82095427105 Mon Sep 17 00:00:00 2001 From: Andrew Date: Wed, 24 Apr 2024 14:00:48 -0400 Subject: [PATCH] Gamma-ray packet source refactor (#2546) * First steps Radioactive packet source GXPacket collection Some initial methods * Add more methods to the packet source * Documentation, cleanup, positronium handling * Black formatting, cleanup * Fix kappa_calculation import loop * Working packet source Creates a Packet collection with all necessary packet properties * Positron deposition and packet tracker arrays * Fix parents dict * Result now mostly consistent with master * Fix kappa calculation test * Adds missing docstring for calculate_energy_factors * Fix util tests and black * Testing skeleton for packet source * black tests * Address comments * Beginnings of an alternate packet source for the dataframe * Possible deposition fix, import fixes * Added old functions * Cleanup, comments, and more dataframe sampling setup * Corrected packet time computation * Modified sampling to use positron information * Remove packet sampling call and replace with dataframe method * Update method call to match current usage * Comment responses * More comments, improve packet dataframe index handling --------- Co-authored-by: Knights-Templars --- tardis/energy_input/GXPacket.py | 30 + tardis/energy_input/__init__.py | 2 - tardis/energy_input/gamma_packet_loop.py | 7 +- tardis/energy_input/gamma_ray_estimators.py | 2 +- tardis/energy_input/gamma_ray_interactions.py | 7 +- .../energy_input/gamma_ray_packet_source.py | 855 ++++++++++++++++++ tardis/energy_input/gamma_ray_transport.py | 245 ++--- tardis/energy_input/main_gamma_ray_loop.py | 127 ++- tardis/energy_input/samplers.py | 28 +- .../tests/test_gamma_ray_packet_source.py | 77 ++ tardis/energy_input/tests/test_util.py | 40 +- tardis/energy_input/util.py | 20 +- .../montecarlo/montecarlo_numba/opacities.py | 21 +- .../montecarlo_numba/tests/test_opacities.py | 38 +- 14 files changed, 1171 insertions(+), 328 deletions(-) create mode 100644 tardis/energy_input/gamma_ray_packet_source.py create mode 100644 tardis/energy_input/tests/test_gamma_ray_packet_source.py diff --git a/tardis/energy_input/GXPacket.py b/tardis/energy_input/GXPacket.py index 5ee44d60896..dd8792876bc 100644 --- a/tardis/energy_input/GXPacket.py +++ b/tardis/energy_input/GXPacket.py @@ -79,6 +79,36 @@ def get_location_r(self): ) +class GXPacketCollection: + """ + Gamma-ray packet collection + """ + + def __init__( + self, + location, + direction, + energy_rf, + energy_cmf, + nu_rf, + nu_cmf, + status, + shell, + time_current, + ): + self.location = location + self.direction = direction + self.energy_rf = energy_rf + self.energy_cmf = energy_cmf + self.nu_rf = nu_rf + self.nu_cmf = nu_cmf + self.status = status + self.shell = shell + self.time_current = time_current + self.number_of_packets = len(self.energy_rf) + self.tau = -np.log(np.random.random(self.number_of_packets)) + + # @njit(**njit_dict_no_parallel) def initialize_packet_properties( isotope_energy, diff --git a/tardis/energy_input/__init__.py b/tardis/energy_input/__init__.py index 5a770c7c2d6..4211b689089 100644 --- a/tardis/energy_input/__init__.py +++ b/tardis/energy_input/__init__.py @@ -1,5 +1,3 @@ """ Contains classes and functions to handle energy deposition and transport. """ - -from tardis.energy_input.util import * diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py index d910a62a6a5..8f76d1f2c07 100644 --- a/tardis/energy_input/gamma_packet_loop.py +++ b/tardis/energy_input/gamma_packet_loop.py @@ -7,6 +7,7 @@ photoabsorption_opacity_calculation, pair_creation_opacity_calculation, photoabsorption_opacity_calculation_kasen, + kappa_calculation, pair_creation_opacity_artis, SIGMA_T, ) @@ -18,7 +19,6 @@ doppler_factor_3d, C_CGS, H_CGS_KEV, - kappa_calculation, get_index, ) from tardis.energy_input.GXPacket import GXPacketStatus @@ -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_estimators.py b/tardis/energy_input/gamma_ray_estimators.py index a17d584225f..c1df7e475dd 100644 --- a/tardis/energy_input/gamma_ray_estimators.py +++ b/tardis/energy_input/gamma_ray_estimators.py @@ -6,13 +6,13 @@ compton_opacity_calculation, SIGMA_T, photoabsorption_opacity_calculation, + kappa_calculation, ) from tardis.energy_input.util import ( angle_aberration_gamma, doppler_factor_3d, H_CGS_KEV, ELECTRON_MASS_ENERGY_KEV, - kappa_calculation, ) diff --git a/tardis/energy_input/gamma_ray_interactions.py b/tardis/energy_input/gamma_ray_interactions.py index 131738c5aab..603601be16f 100644 --- a/tardis/energy_input/gamma_ray_interactions.py +++ b/tardis/energy_input/gamma_ray_interactions.py @@ -2,10 +2,12 @@ from numba import njit from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel -from tardis.montecarlo.montecarlo_numba.opacities import compton_opacity_partial +from tardis.montecarlo.montecarlo_numba.opacities import ( + compton_opacity_partial, + kappa_calculation, +) from tardis.energy_input.util import ( get_random_unit_vector, - kappa_calculation, euler_rodrigues, compton_theta_distribution, get_perpendicular_vector, @@ -157,7 +159,6 @@ def get_compton_fraction_urilight(energy): accept = False while not accept: - z = np.random.random(3) alpha1 = np.log(1.0 / x0) alpha2 = (1.0 - x0**2.0) / 2.0 diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py new file mode 100644 index 00000000000..a5d1522043e --- /dev/null +++ b/tardis/energy_input/gamma_ray_packet_source.py @@ -0,0 +1,855 @@ +import numpy as np +import pandas as pd + +from tardis.energy_input.energy_source import ( + positronium_continuum, +) +from tardis.energy_input.GXPacket import ( + GXPacketCollection, +) +from tardis.energy_input.samplers import sample_energy +from tardis.energy_input.util import ( + H_CGS_KEV, + doppler_factor_3d, + get_index, + get_random_unit_vector, +) +from tardis.montecarlo.packet_source import BasePacketSource + + +class RadioactivePacketSource(BasePacketSource): + def __init__( + self, + packet_energy, + gamma_ray_lines, + positronium_fraction, + inner_velocities, + outer_velocities, + inv_volume_time, + times, + energy_df_rows, + effective_times, + taus, + parents, + average_positron_energies, + average_power_per_mass, + **kwargs, + ): + self.packet_energy = packet_energy + self.gamma_ray_lines = gamma_ray_lines + self.positronium_fraction = positronium_fraction + self.inner_velocities = inner_velocities + self.outer_velocities = outer_velocities + self.inv_volume_time = inv_volume_time + self.times = times + self.energy_df_rows = energy_df_rows + self.effective_times = effective_times + self.taus = taus + 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): + return super().create_packet_mus(no_of_packets, *args, **kwargs) + + def create_packet_radii( + self, no_of_packets, inner_velocity, outer_velocity + ): + """Initialize the random radii of packets in a shell + + Parameters + ---------- + packet_count : int + Number of packets in the shell + inner_velocity : float + Inner velocity of the shell + outer_velocity : float + Outer velocity of the shell + + Returns + ------- + array + Array of length packet_count of random locations in the shell + """ + z = np.random.random(no_of_packets) + initial_radii = ( + z * inner_velocity**3.0 + (1.0 - z) * outer_velocity**3.0 + ) ** (1.0 / 3.0) + + return initial_radii + + def create_packet_nus( + self, + no_of_packets, + energy, + intensity, + positronium_fraction, + positronium_energy, + positronium_intensity, + ): + """Create an array of packet frequency-energies (i.e. E = h * nu) + + Parameters + ---------- + no_of_packets : int + Number of packets to produce frequency-energies for + energy : One-dimensional Numpy Array, dtype float + Array of frequency-energies to sample + intensity : One-dimensional Numpy Array, dtype float + Array of intensities to sample + positronium_fraction : float + The fraction of positrons that form positronium + positronium_energy : array + Array of positronium frequency-energies to sample + positronium_intensity : array + Array of positronium intensities to sample + + Returns + ------- + array + Array of sampled frequency-energies + array + Positron creation mask + """ + nu_energies = np.zeros(no_of_packets) + positrons = np.zeros(no_of_packets) + zs = np.random.random(no_of_packets) + for i in range(no_of_packets): + nu_energies[i] = sample_energy(energy, intensity) + # positron + if nu_energies[i] == 511: + # positronium formation 25% of the time if fraction is 1 + if zs[i] < positronium_fraction and np.random.random() < 0.25: + nu_energies[i] = sample_energy( + positronium_energy, positronium_intensity + ) + positrons[i] = 1 + + return nu_energies, positrons + + def create_packet_directions(self, no_of_packets): + """Create an array of random directions + + Parameters + ---------- + no_of_packets : int + Number of packets to produce directions for + + Returns + ------- + array + Array of direction vectors + """ + directions = np.zeros((3, no_of_packets)) + for i in range(no_of_packets): + directions[:, i] = get_random_unit_vector() + + return directions + + def create_packet_energies(self, no_of_packets, energy): + """Create the uniform packet energy for a number of packets + + Parameters + ---------- + no_of_packets : int + Number of packets + energy : float + The packet energy + + Returns + ------- + array + Array of packet energies + """ + return np.ones(no_of_packets) * energy + + def create_packet_times_uniform_time(self, no_of_packets, start, end): + """Samples decay time uniformly (needs non-uniform packet energies) + + Parameters + ---------- + no_of_packets : int + Number of packets + start : float + Start time + end : float + End time + + Returns + ------- + array + Array of packet decay times + """ + z = np.random.random(no_of_packets) + decay_times = z * start + (1 - z) * end + return decay_times + + def create_packet_times_uniform_energy( + self, + no_of_packets, + start_tau, + end_tau=0.0, + decay_time_min=0.0, + decay_time_max=0.0, + ): + """Samples the decay time from the mean lifetime of the isotopes + + Parameters + ---------- + no_of_packets : int + Number of packets + start_tau : float + Initial isotope mean lifetime + end_tau : float, optional + Ending mean lifetime, by default 0.0 for single decays + decay_time_min : float, optional + Minimum time to decay, by default 0.0 + decay_time_max : float, optional + Maximum time to decay, by default 0.0 + + Returns + ------- + array + Array of decay times + """ + decay_times = np.ones(no_of_packets) * decay_time_min + for i in range(no_of_packets): + # rejection sampling + while (decay_times[i] <= decay_time_min) or ( + decay_times[i] >= decay_time_max + ): + decay_times[i] = -start_tau * np.log( + np.random.random() + ) - end_tau * np.log(np.random.random()) + return decay_times + + def calculate_energy_factors(self, no_of_packets, start_time, decay_times): + """Calculates the factors that adjust the energy of packets emitted + before the first time step and moves those packets to the earliest + possible time + + Parameters + ---------- + no_of_packets : int + Number of packets + start_time : float + First time step + decay_times : array + Packet decay times + + Returns + ------- + array + Energy factors + array + Adjusted decay times + """ + energy_factors = np.ones(no_of_packets) + for i in range(no_of_packets): + if decay_times[i] < start_time: + energy_factors[i] = decay_times[i] / start_time + decay_times[i] = start_time + return energy_factors, decay_times + + def create_packets(self, decays_per_isotope, *args, **kwargs): + """Initialize a collection of GXPacket objects for the simulation + to operate on. + + Parameters + ---------- + decays_per_isotope : array int64 + Number of decays per simulation shell per isotope + + Returns + ------- + list + List of GXPacket objects + array + Array of main output dataframe rows + array + Array of plotting output dataframe rows + array + Array of positron output dataframe rows + """ + number_of_packets = decays_per_isotope.sum().sum() + decays_per_shell = decays_per_isotope.sum().values + + locations = np.zeros((3, number_of_packets)) + directions = np.zeros((3, number_of_packets)) + packet_energies_rf = np.zeros(number_of_packets) + packet_energies_cmf = np.zeros(number_of_packets) + nus_rf = np.zeros(number_of_packets) + nus_cmf = np.zeros(number_of_packets) + shells = np.zeros(number_of_packets) + times = np.zeros(number_of_packets) + # set packets to IN_PROCESS status + statuses = np.ones(number_of_packets, dtype=np.int64) * 3 + + 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): + isotope_packet_count_df = decays_per_isotope.T.iloc[shell_number] + + for isotope_name, isotope_packet_count in zip( + self.gamma_ray_lines.keys(), isotope_packet_count_df.values + ): + isotope_energy = self.gamma_ray_lines[isotope_name][0, :] + isotope_intensity = self.gamma_ray_lines[isotope_name][1, :] + isotope_positron_fraction = self.calculate_positron_fraction( + self.average_positron_energies[isotope_name], + isotope_energy, + isotope_intensity, + ) + tau_start = self.taus[isotope_name] + + if isotope_name in self.parents: + tau_end = self.taus[self.parents[isotope_name]] + else: + tau_end = 0 + + # sample radii at time = 0 + initial_radii = self.create_packet_radii( + isotope_packet_count, + self.inner_velocities[shell_number], + self.outer_velocities[shell_number], + ) + + # sample directions (valid at all times) + initial_directions = self.create_packet_directions( + isotope_packet_count + ) + + # packet decay time + initial_times = self.create_packet_times_uniform_energy( + isotope_packet_count, + tau_start, + tau_end, + decay_time_min=0, + decay_time_max=self.times[-1], + ) + + # get the time step index of the packets + initial_time_indexes = np.array( + [ + get_index(decay_time, self.times) + for decay_time in initial_times + ] + ) + + # get the time of the middle of the step for each packet + packet_effective_times = np.array( + [self.effective_times[i] for i in initial_time_indexes] + ) + + # 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. + initial_locations = ( + initial_radii + * packet_effective_times + * self.create_packet_directions(isotope_packet_count) + ) + + # get the packet shell index + initial_shells = np.ones(isotope_packet_count) * shell_number + + # the individual gamma-ray energies that make up a packet + # co-moving frame, including positronium formation + initial_nu_energies_cmf, positron_mask = self.create_packet_nus( + isotope_packet_count, + isotope_energy, + isotope_intensity, + self.positronium_fraction, + positronium_energy, + positronium_intensity, + ) + + # equivalent frequencies + initial_nus_cmf = initial_nu_energies_cmf / H_CGS_KEV + + # compute scaling factor for packets emitted before start time + # and move packets to start at that time + # probably not necessary- we have rejection sampling in the + # create_packet_times_uniform_energy method + energy_factors, initial_times = self.calculate_energy_factors( + isotope_packet_count, self.times[0], initial_times + ) + + # the CMF energy of a packet scaled by the "early energy factor" + initial_packet_energies_cmf = ( + self.create_packet_energies( + isotope_packet_count, self.packet_energy + ) + * energy_factors + ) + + # rest frame gamma-ray energy and frequency + # this probably works fine without the loop + initial_packet_energies_rf = np.zeros(isotope_packet_count) + initial_nus_rf = np.zeros(isotope_packet_count) + for i in range(isotope_packet_count): + doppler_factor = doppler_factor_3d( + initial_directions[:, i], + initial_locations[:, i], + initial_times[i], + ) + initial_packet_energies_rf[i] = ( + 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, + # * 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 + ) + + # collect packet properties + locations[ + :, packet_index - isotope_packet_count : packet_index + ] = initial_locations + directions[ + :, packet_index - isotope_packet_count : packet_index + ] = initial_directions + packet_energies_rf[ + packet_index - isotope_packet_count : packet_index + ] = initial_packet_energies_rf + packet_energies_cmf[ + packet_index - isotope_packet_count : packet_index + ] = initial_packet_energies_cmf + nus_rf[ + packet_index - isotope_packet_count : packet_index + ] = initial_nus_rf + nus_cmf[ + packet_index - isotope_packet_count : packet_index + ] = initial_nus_cmf + shells[ + packet_index - isotope_packet_count : packet_index + ] = initial_shells + times[ + packet_index - isotope_packet_count : packet_index + ] = initial_times + + return GXPacketCollection( + locations, + directions, + packet_energies_rf, + packet_energies_cmf, + nus_rf, + nus_cmf, + statuses, + shells, + times, + ) + + def calculate_positron_fraction( + self, positron_energy, isotope_energy, isotope_intensity + ): + """Calculate the fraction of energy that an isotope + releases as positron kinetic energy + + Parameters + ---------- + positron_energy : float + Average kinetic energy of positrons from decay + isotope_energy : numpy array + Photon energies released by the isotope + isotope_intensity : numpy array + Intensity of photon energy release + + Returns + ------- + float + Fraction of energy released as positron kinetic energy + """ + return positron_energy / np.sum(isotope_energy * isotope_intensity) + + +class GammaRayPacketSource(BasePacketSource): + def __init__( + self, + packet_energy, + gamma_ray_lines, + positronium_fraction, + inner_velocities, + outer_velocities, + inv_volume_time, + times, + energy_df_rows, + effective_times, + taus, + parents, + average_positron_energies, + average_power_per_mass, + **kwargs, + ): + self.packet_energy = packet_energy + self.gamma_ray_lines = gamma_ray_lines + self.positronium_fraction = positronium_fraction + self.inner_velocities = inner_velocities + self.outer_velocities = outer_velocities + self.inv_volume_time = inv_volume_time + self.times = times + self.energy_df_rows = energy_df_rows + self.effective_times = effective_times + self.taus = taus + 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): + return super().create_packet_mus(no_of_packets, *args, **kwargs) + + def create_packet_radii(self, sampled_packets_df): + """Initialize the random radii of packets in a shell + + Parameters + ---------- + packet_count : int + Number of packets in the shell + sampled_packets_df : pd.DataFrame + Dataframe where each row is a packet + + Returns + ------- + array + Array of length packet_count of random locations in the shell + """ + z = np.random.random(len(sampled_packets_df)) + initial_radii = ( + z * sampled_packets_df["inner_velocity"] ** 3.0 + + (1.0 - z) * sampled_packets_df["outer_velocity"] ** 3.0 + ) ** (1.0 / 3.0) + + return initial_radii + + def create_packet_nus( + self, + no_of_packets, + packets, + positronium_fraction, + positronium_energy, + positronium_intensity, + ): + """Create an array of packet frequency-energies (i.e. E = h * nu) + + Parameters + ---------- + no_of_packets : int + Number of packets to produce frequency-energies for + packets : pd.DataFrame + DataFrame of packets + positronium_fraction : float + The fraction of positrons that form positronium + positronium_energy : array + Array of positronium frequency-energies to sample + positronium_intensity : array + Array of positronium intensities to sample + + Returns + ------- + array + Array of sampled frequency-energies + """ + energy_array = np.zeros(no_of_packets) + zs = np.random.random(no_of_packets) + for i in range(no_of_packets): + # positron + if packets.iloc[i]["decay_type"] == "bp": + # positronium formation 75% of the time if fraction is 1 + if zs[i] < positronium_fraction and np.random.random() < 0.75: + energy_array[i] = sample_energy( + positronium_energy, positronium_intensity + ) + else: + energy_array[i] = 511 + else: + energy_array[i] = packets.iloc[i]["radiation_energy_kev"] + + return energy_array + + def create_packet_directions(self, no_of_packets): + """Create an array of random directions + + Parameters + ---------- + no_of_packets : int + Number of packets to produce directions for + + Returns + ------- + array + Array of direction vectors + """ + directions = np.zeros((3, no_of_packets)) + for i in range(no_of_packets): + directions[:, i] = get_random_unit_vector() + + return directions + + def create_packet_energies(self, no_of_packets, energy): + """Create the uniform packet energy for a number of packets + + Parameters + ---------- + no_of_packets : int + Number of packets + energy : float + The packet energy + + Returns + ------- + array + Array of packet energies + """ + return np.ones(no_of_packets) * energy + + def create_packet_times_uniform_time(self, no_of_packets, start, end): + """Samples decay time uniformly (needs non-uniform packet energies) + + Parameters + ---------- + no_of_packets : int + Number of packets + start : float + Start time + end : float + End time + + Returns + ------- + array + Array of packet decay times + """ + z = np.random.random(no_of_packets) + decay_times = z * start + (1 - z) * end + return decay_times + + def create_packet_times_uniform_energy( + self, no_of_packets, isotopes, decay_time + ): + """Samples the decay time from the mean lifetime of the isotopes + + Parameters + ---------- + no_of_packets : int + Number of packets + isotopes : pd.Series + Series of packet parent isotopes + decay_time : array + Series of packet decay time index + + Returns + ------- + array + Array of decay times + """ + decay_times = np.zeros(len(no_of_packets)) + for i, isotope in enumerate(isotopes.to_numpy()): + decay_time_min = self.times[decay_time[i]] + if decay_time_min == self.times[-1]: + decay_time_max = self.effective_times[-1] + else: + decay_time_max = self.times[decay_time[i] + 1] + # rejection sampling + while (decay_times[i] <= decay_time_min) or ( + decay_times[i] >= decay_time_max + ): + decay_times[i] = -self.taus[isotope] * np.log( + np.random.random() + ) + return decay_times + + def create_packets( + self, decays_per_isotope, number_of_packets, *args, **kwargs + ): + """Initialize a collection of GXPacket objects for the simulation + to operate on. + + Parameters + ---------- + decays_per_isotope : array int64 + Probability of decays per simulation shell per isotope per time step + number_of_packets : int + Number of packets to create + + Returns + ------- + GXPacketCollection + """ + # initialize arrays for most packet properties + locations = np.zeros((3, number_of_packets)) + directions = np.zeros((3, number_of_packets)) + packet_energies_rf = np.zeros(number_of_packets) + packet_energies_cmf = np.zeros(number_of_packets) + nus_rf = np.zeros(number_of_packets) + nus_cmf = np.zeros(number_of_packets) + times = np.zeros(number_of_packets) + # set packets to IN_PROCESS status + statuses = np.ones(number_of_packets, dtype=np.int64) * 3 + + self.energy_plot_positron_rows = np.zeros((number_of_packets, 4)) + + # compute positronium continuum + positronium_energy, positronium_intensity = positronium_continuum() + + # sample packets from dataframe, returning a dataframe where each row is + # a sampled packet + sampled_packets_df = decays_per_isotope.sample( + n=number_of_packets, + weights="decay_energy_erg", + replace=True, + random_state=np.random.RandomState(self.base_seed), + ) + # get unique isotopes that have produced packets + isotopes = pd.unique(sampled_packets_df.index.get_level_values(2)) + + # compute the positron fraction for unique isotopes + isotope_positron_fraction = self.calculate_positron_fraction(isotopes) + + # get the packet shell index + shells = sampled_packets_df.index.get_level_values(1) + + # get the inner and outer velocity boundaries for each packet to compute + # the initial radii + sampled_packets_df["inner_velocity"] = self.inner_velocities[shells] + sampled_packets_df["outer_velocity"] = self.outer_velocities[shells] + + # sample radii at time = 0 + initial_radii = self.create_packet_radii(sampled_packets_df) + + # get the time step index of the packets + initial_time_indexes = sampled_packets_df.index.get_level_values(0) + + # get the time of the middle of the step for each packet + packet_effective_times = np.array( + [self.effective_times[i] for i in initial_time_indexes] + ) + + # packet decay time + times = self.create_packet_times_uniform_energy( + number_of_packets, + sampled_packets_df.index.get_level_values(2), + packet_effective_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 + * packet_effective_times + * self.create_packet_directions(number_of_packets) + ) + + # sample directions (valid at all times), non-relativistic + directions = self.create_packet_directions(number_of_packets) + + # the individual gamma-ray energy that makes up a packet + # co-moving frame, including positronium formation + nu_energies_cmf = self.create_packet_nus( + number_of_packets, + sampled_packets_df, + self.positronium_fraction, + positronium_energy, + positronium_intensity, + ) + + # equivalent frequencies + nus_cmf = nu_energies_cmf / H_CGS_KEV + + # per packet co-moving frame total energy + packet_energies_cmf = self.create_packet_energies( + number_of_packets, self.packet_energy + ) + + # rest frame gamma-ray energy and frequency + # this probably works fine without the loop + # non-relativistic + packet_energies_rf = np.zeros(number_of_packets) + nus_rf = np.zeros(number_of_packets) + for i in range(number_of_packets): + doppler_factor = doppler_factor_3d( + directions[:, i], + locations[:, i], + times[i], + ) + packet_energies_rf[i] = packet_energies_cmf[i] / doppler_factor + nus_rf[i] = nus_cmf[i] / doppler_factor + + # deposit positron energy in both output arrays + # this is an average across all packets that are created + # it could be changed to be only for packets that are from positrons + self.energy_plot_positron_rows[i] = np.array( + [ + i, + isotope_positron_fraction[sampled_packets_df["isotopes"][i]] + * packet_energies_cmf[i], + # this needs to be sqrt(sum of squares) to get radius + np.linalg.norm(locations[i]), + times[i], + ] + ) + + # this is an average across all packets that are created + # it could be changed to be only for packets that are from positrons + self.energy_df_rows[shells[i], times[i]] += ( + isotope_positron_fraction[sampled_packets_df["isotopes"][i]] + * packet_energies_cmf[i] + ) + + return GXPacketCollection( + locations, + directions, + packet_energies_rf, + packet_energies_cmf, + nus_rf, + nus_cmf, + statuses, + shells, + times, + ) + + def calculate_positron_fraction(self, isotopes): + """Calculate the fraction of energy that an isotope + releases as positron kinetic energy + + Parameters + ---------- + isotopes : array + Array of isotope names as strings + + Returns + ------- + dict + Fraction of energy released as positron kinetic energy per isotope + """ + positron_fraction = {} + + for isotope in isotopes: + isotope_energy = self.gamma_ray_lines[isotope][0, :] + isotope_intensity = self.gamma_ray_lines[isotope][1, :] + positron_fraction[isotope] = self.average_positron_energies[ + isotope + ] / np.sum(isotope_energy * isotope_intensity) + return positron_fraction diff --git a/tardis/energy_input/gamma_ray_transport.py b/tardis/energy_input/gamma_ray_transport.py index efaaad6a6b0..0412195e1e0 100644 --- a/tardis/energy_input/gamma_ray_transport.py +++ b/tardis/energy_input/gamma_ray_transport.py @@ -3,17 +3,11 @@ import pandas as pd import astropy.units as u import radioactivedecay as rd -from numba import njit -from numba.typed import List from tardis.energy_input.energy_source import ( get_all_isotopes, - positronium_continuum, setup_input_energy, ) -from tardis.energy_input.GXPacket import initialize_packet_properties -from tardis.energy_input.samplers import initial_packet_radius -from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel from tardis.montecarlo.montecarlo_numba.opacities import M_P # Energy: keV, exported as eV for SF solver @@ -84,189 +78,103 @@ def get_chain_decay_power_per_ejectamass( return decaypower -@njit(**njit_dict_no_parallel) -def calculate_positron_fraction( - positron_energy, isotope_energy, isotope_intensity -): - """Calculate the fraction of energy that an isotope - releases as positron kinetic energy +def calculate_ejecta_velocity_volume(model): + outer_velocities = model.v_outer.to("cm/s").value + inner_velocities = model.v_inner.to("cm/s").value + ejecta_velocity_volume = ( + 4 * np.pi / 3 * (outer_velocities**3.0 - inner_velocities**3.0) + ) + + return ejecta_velocity_volume + + +def calculate_total_decays_old(inventories, time_delta): + """Function to create inventories of isotope Parameters ---------- - positron_energy : float - Average kinetic energy of positrons from decay - isotope_energy : numpy array - Photon energies released by the isotope - isotope_intensity : numpy array - Intensity of photon energy release + model : tardis.Radial1DModel + The tardis model to calculate gamma ray propagation through + + time_end : float + End time of simulation in days Returns ------- - float - Fraction of energy released as positron kinetic energy + Total decay list : List + list of total decays for x g of isotope for time 't' """ - return positron_energy / np.sum(isotope_energy * isotope_intensity) + time_delta = u.Quantity(time_delta, u.s) + total_decays = {} + for shell, isotopes in inventories.items(): + total_decays[shell] = {} + for isotope, name in isotopes.items(): + # decays = name.decay(time_delta.value, "s") + total_decays[shell][isotope] = name.cumulative_decays( + time_delta.value + ) + return total_decays -def initialize_packets( - decays_per_isotope, - packet_energy, - positronium_fraction, - inner_velocities, - outer_velocities, - gamma_ray_lines, - average_positron_energies, - inv_volume_time, - times, - energy_df_rows, - effective_times, - taus, - parents, - average_power_per_mass, -): - """Initialize a list of GXPacket objects for the simulation - to operate on. +def create_isotope_dicts_old(raw_isotope_abundance, cell_masses): + """ + Function to create a dictionary of isotopes for each shell with their masses. Parameters ---------- - decays_per_isotope : array int64 - Number of decays per simulation shell per isotope - input_energy : float64 - Total input energy from decay - ni56_lines : array float64 - Lines and intensities for Ni56 - co56_lines : array float64 - Lines and intensities for Co56 - inner_velocities : array float64 - Inner velocities of the shells - outer_velocities : array float64 - Outer velocities of the shells - inv_volume_time : array float64 - Inverse volume with time - times : array float64 - Simulation time steps - energy_df_rows : list - Setup list for energy DataFrame output - effective_times : array float64 - Middle time of the time step - taus : array float64 - Mean lifetime for each isotope + raw_isotope_abundance : pd.DataFrame + isotope abundance in mass fractions. + cell_masses : numpy.ndarray + shell masses in units of g Returns ------- - list - List of GXPacket objects - array - Array of main output dataframe rows - array - Array of plotting output dataframe rows - array - Array of positron output dataframe rows - """ - packets = List() - - number_of_packets = decays_per_isotope.sum().sum() - decays_per_shell = decays_per_isotope.sum().values - - energy_plot_df_rows = np.zeros((number_of_packets, 8)) - energy_plot_positron_rows = np.zeros((number_of_packets, 4)) - - positronium_energy, positronium_intensity = positronium_continuum() - isotopes = list(gamma_ray_lines.keys()) + isotope_dicts : Dict + dictionary of isotopes for each shell with their ``masses``. + For eg: {0: {(28, 56): {'Ni56': 0.0001}, (27, 57): {'Co56': 0.0001}} + {1: {(28, 56): {'Ni56': 0.0001}, (27, 57): {'Co56': 0.0001}}} etc - packet_index = 0 - logger.info("Isotope packet count dataframe") - for shell_number, pkts in enumerate(decays_per_shell): - initial_radii = initial_packet_radius( - pkts, inner_velocities[shell_number], outer_velocities[shell_number] - ) - - isotope_packet_count_df = decays_per_isotope.T.iloc[shell_number] - logger.info(isotope_packet_count_df) - i = 0 - for isotope_name, isotope_packet_count in zip( - isotopes, isotope_packet_count_df.values - ): - isotope_energy = gamma_ray_lines[isotope_name][0, :] - isotope_intensity = gamma_ray_lines[isotope_name][1, :] - isotope = rd.Nuclide(isotope_name) - atomic_number = isotope.Z - mass_number = isotope.A - isotope_positron_fraction = calculate_positron_fraction( - average_positron_energies[isotope_name], - isotope_energy, - isotope_intensity, + """ + 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] * cell_masses[i].to(u.g).value ) - tau_start = taus[isotope_name] - - if isotope_name in parents: - tau_end = taus[parents[isotope_name]] - else: - tau_end = 0 - - for c in range(isotope_packet_count): - packet, decay_time_index = initialize_packet_properties( - isotope_energy, - isotope_intensity, - positronium_energy, - positronium_intensity, - positronium_fraction, - packet_energy, - shell_number, - tau_start, - tau_end, - initial_radii[c], - times, - effective_times, - average_power_per_mass, - ) - - energy_df_rows[shell_number, decay_time_index] += ( - isotope_positron_fraction * packet_energy * 1000 - ) - energy_plot_df_rows[packet_index] = np.array( - [ - i, - packet.energy_rf, - packet.get_location_r(), - packet.time_current, - int(packet.status), - 0, - 0, - 0, - ] - ) - - energy_plot_positron_rows[packet_index] = [ - packet_index, - isotope_positron_fraction * packet_energy * 1000, - # * inv_volume_time[packet.shell, decay_time_index], - packet.get_location_r(), - packet.time_current, - ] + return isotope_dicts - packets.append(packet) - i += 1 - packet_index += 1 - - return ( - packets, - energy_df_rows, - energy_plot_df_rows, - energy_plot_positron_rows, - ) +def create_inventories_dict_old(isotope_dict): + """Function to create dictionary of inventories for each shell + Parameters + ---------- + isotope_dict : Dict + dictionary of isotopes for each shell with their ``masses``. -def calculate_ejecta_velocity_volume(model): - outer_velocities = model.v_outer.to("cm/s").value - inner_velocities = model.v_inner.to("cm/s").value - ejecta_velocity_volume = ( - 4 * np.pi / 3 * (outer_velocities**3.0 - inner_velocities**3.0) - ) + Returns + ------- + inv : Dict + dictionary of inventories for each shell + For eg: {0: {'Ni56': , + 'Co56': }, + {1: {'Ni56': , + 'Co56': }} etc + """ + inv = {} + for shell, isotopes in isotope_dict.items(): + inv[shell] = {} + for isotope, name in isotopes.items(): + inv[shell][isotope] = rd.Inventory(name, "g") - return ejecta_velocity_volume + return inv def calculate_average_energies(raw_isotope_abundance, gamma_ray_lines): @@ -362,7 +270,10 @@ def get_taus(raw_isotope_abundance): if child is not None: for c in child: if rd.Nuclide(c).half_life("readable") != "stable": - parents[isotope] = c + # this is a dict of child: parent intended to find + # the parents of a given isotope. + # if there is no parent, there is no item. + parents[c] = isotope return taus, parents diff --git a/tardis/energy_input/main_gamma_ray_loop.py b/tardis/energy_input/main_gamma_ray_loop.py index 1c82cb48e00..654656300ad 100644 --- a/tardis/energy_input/main_gamma_ray_loop.py +++ b/tardis/energy_input/main_gamma_ray_loop.py @@ -1,30 +1,36 @@ import logging + +import astropy.units as u import numpy as np import pandas as pd -import astropy.units as u from tardis.energy_input.energy_source import ( - get_all_isotopes, get_nuclear_lines_database, ) - from tardis.energy_input.gamma_packet_loop import gamma_packet_loop -from tardis.energy_input.gamma_ray_transport import ( - initialize_packets, - calculate_ejecta_velocity_volume, - create_isotope_dicts, - create_inventories_dict, +from tardis.energy_input.gamma_ray_channel import ( calculate_total_decays, + create_inventories_dict, + create_isotope_dicts, +) + +from tardis.energy_input.gamma_ray_transport import ( + calculate_total_decays_old, + create_isotope_dicts_old, + create_inventories_dict_old, +) +from tardis.energy_input.gamma_ray_packet_source import RadioactivePacketSource +from tardis.energy_input.gamma_ray_transport import ( calculate_average_energies, - decay_chain_energies, - calculate_energy_per_mass, calculate_average_power_per_mass, - iron_group_fraction_per_shell, - get_taus, - fractional_decay_energy, - packets_per_isotope, + calculate_ejecta_velocity_volume, + calculate_energy_per_mass, + decay_chain_energies, distribute_packets, + get_taus, + iron_group_fraction_per_shell, ) +from tardis.energy_input.GXPacket import GXPacket logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) @@ -104,9 +110,9 @@ def run_gamma_ray_loop( 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( + isotope_dict = create_isotope_dicts_old(raw_isotope_abundance, shell_masses) + inventories_dict = create_inventories_dict_old(isotope_dict) + total_decays = calculate_total_decays_old( inventories_dict, time_end - time_start ) @@ -132,17 +138,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 = ( @@ -164,49 +159,69 @@ def run_gamma_ray_loop( logger.info("Initializing packets") - ( - packets, - energy_df_rows, - energy_plot_df_rows, - energy_plot_positron_rows, - ) = initialize_packets( - packets_per_isotope_df, + packet_source = RadioactivePacketSource( individual_packet_energy, + gamma_ray_line_dict, positronium_fraction, inner_velocities, outer_velocities, - gamma_ray_line_dict, - average_positron_energies, inv_volume_time, times, energy_df_rows, effective_time_array, taus, parents, + average_positron_energies, average_power_per_mass, ) - 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 + 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 = packet_collection.energy_cmf.sum() + total_rf_energy = packet_collection.energy_rf.sum() + for i in range(number_of_packets): + packet = GXPacket( + packet_collection.location[:, i], + packet_collection.direction[:, i], + packet_collection.energy_rf[i], + packet_collection.energy_cmf[i], + packet_collection.nu_rf[i], + packet_collection.nu_cmf[i], + packet_collection.status[i], + packet_collection.shell[i], + packet_collection.time_current[i], + ) + packets.append(packet) + energy_plot_df_rows[i] = np.array( + [ + i, + packet.energy_rf, + packet.get_location_r(), + packet.time_current, + int(packet.status), + 0, + 0, + 0, + ] + ) 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( @@ -227,7 +242,6 @@ def run_gamma_ray_loop( energy_df_rows, energy_plot_df_rows, energy_out, - packets_out, packets_info_array, ) @@ -246,7 +260,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", @@ -255,22 +269,6 @@ 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 @@ -287,5 +285,4 @@ def run_gamma_ray_loop( decayed_packet_count, energy_plot_positrons, energy_estimated_deposition, - packets_out_df, ) diff --git a/tardis/energy_input/samplers.py b/tardis/energy_input/samplers.py index 15ec2186ae3..aa91a665f6c 100644 --- a/tardis/energy_input/samplers.py +++ b/tardis/energy_input/samplers.py @@ -105,7 +105,7 @@ def sample_energy(energy, intensity): average = (energy * intensity).sum() total = 0 - for (e, i) in zip(energy, intensity): + for e, i in zip(energy, intensity): total += e * i / average if z <= total: return e @@ -138,29 +138,3 @@ def sample_decay_time( np.random.random() ) return decay_time - - -@njit(**njit_dict_no_parallel) -def initial_packet_radius(packet_count, inner_velocity, outer_velocity): - """Initialize the random radii of packets in a shell - - Parameters - ---------- - packet_count : int - Number of packets in the shell - inner_velocity : float - Inner velocity of the shell - outer_velocity : float - Outer velocity of the shell - - Returns - ------- - array - Array of length packet_count of random locations in the shell - """ - z = np.random.random(packet_count) - initial_radii = ( - z * inner_velocity**3.0 + (1.0 - z) * outer_velocity**3.0 - ) ** (1.0 / 3.0) - - return initial_radii diff --git a/tardis/energy_input/tests/test_gamma_ray_packet_source.py b/tardis/energy_input/tests/test_gamma_ray_packet_source.py new file mode 100644 index 00000000000..ff5fd9e8012 --- /dev/null +++ b/tardis/energy_input/tests/test_gamma_ray_packet_source.py @@ -0,0 +1,77 @@ +import numpy as np +import pytest + +from tardis.energy_input.gamma_ray_packet_source import RadioactivePacketSource + + +@pytest.mark.skip(reason="Packet source init is very complex") +class TestGammaRayPacketSource: + @pytest.fixture(scope="class") + def radioactivepacketsource(self, request): + """ + Create RadioactivePacketSource instance. + + Yields + ------- + tardis.energy_input.gamma_ray_packet_source.RadioactivePacketSource + """ + cls = type(self) + cls.packet_source = RadioactivePacketSource(base_seed=1963) + yield cls.packet_source + + def test_create_packet_radii( + self, regression_data, radioactivepacketsource + ): + actual = self.packet_source.create_packet_radii() + expected = regression_data.sync_ndarray(actual) + assert np.all(np.isclose(actual, expected)) + + def test_create_packet_nus(self, regression_data, radioactivepacketsource): + actual = self.packet_source.create_packet_nus() + expected = regression_data.sync_ndarray(actual) + assert np.all(np.isclose(actual, expected)) + + def test_create_packet_directions( + self, regression_data, radioactivepacketsource + ): + actual = self.packet_source.create_packet_directions() + expected = regression_data.sync_ndarray(actual) + assert np.all(np.isclose(actual, expected)) + + def test_create_packet_energies( + self, regression_data, radioactivepacketsource + ): + actual = self.packet_source.create_packet_energies() + expected = regression_data.sync_ndarray(actual) + assert np.all(np.isclose(actual, expected)) + + def test_create_packet_times_uniform_time( + self, regression_data, radioactivepacketsource + ): + actual = self.packet_source.create_packet_times_uniform_time() + expected = regression_data.sync_ndarray(actual) + assert np.all(np.isclose(actual, expected)) + + def test_create_packet_times_uniform_energy( + self, regression_data, radioactivepacketsource + ): + actual = self.packet_source.create_packet_times_uniform_energy() + expected = regression_data.sync_ndarray(actual) + assert np.all(np.isclose(actual, expected)) + + def test_calculate_energy_factors( + self, regression_data, radioactivepacketsource + ): + actual = self.packet_source.calculate_energy_factors() + expected = regression_data.sync_ndarray(actual) + assert np.all(np.isclose(actual, expected)) + + def test_create_packets(self, regression_data, radioactivepacketsource): + assert True + + def test_calculate_positron_fraction( + self, regression_data, radioactivepacketsource + ): + actual = self.packet_source.calculate_positron_fraction() + expected = regression_data.sync_ndarray(actual) + assert np.all(np.isclose(actual, expected)) diff --git a/tardis/energy_input/tests/test_util.py b/tardis/energy_input/tests/test_util.py index 6ea70d8f022..4887f392e99 100644 --- a/tardis/energy_input/tests/test_util.py +++ b/tardis/energy_input/tests/test_util.py @@ -1,15 +1,16 @@ -from random import random -import pytest -import astropy.units as u -import numpy.testing as npt import numpy as np +import numpy.testing as npt +import pytest -import tardis.energy_input.util as util from tardis.energy_input.util import ( R_ELECTRON_SQUARED, get_perpendicular_vector, + klein_nishina, + spherical_to_cartesian, +) +from tardis.montecarlo.montecarlo_numba.opacities import ( + kappa_calculation, ) -from tardis import constants as const @pytest.mark.parametrize( @@ -25,7 +26,7 @@ def test_spherical_to_cartesian( r, theta, phi, expected_x, expected_y, expected_z ): - actual_x, actual_y, actual_z = util.spherical_to_cartesian(r, theta, phi) + actual_x, actual_y, actual_z = spherical_to_cartesian(r, theta, phi) npt.assert_almost_equal(actual_x, expected_x) npt.assert_almost_equal(actual_y, expected_y) npt.assert_almost_equal(actual_z, expected_z) @@ -43,27 +44,6 @@ def test_angle_aberration_gamma(): assert False -@pytest.mark.parametrize( - ["energy", "expected"], - [ - (511.0, 1.0000021334560507), - (255.5, 0.5000010667280254), - (0.0, 0.0), - (511.0e7, 10000021.334560508), - ], -) -def test_kappa_calculation(energy, expected): - """ - - Parameters - ---------- - energy : float - expected : float - """ - kappa = util.kappa_calculation(energy) - npt.assert_almost_equal(kappa, expected) - - @pytest.mark.xfail(reason="To be removed") def test_euler_rodrigues(): """Test Euler-Rodrigues rotation""" @@ -94,9 +74,9 @@ def test_klein_nishina(energy, theta_C): theta_C : float In radians """ - actual = util.klein_nishina(energy, theta_C) + actual = klein_nishina(energy, theta_C) - kappa = util.kappa_calculation(energy) + kappa = kappa_calculation(energy) expected = ( R_ELECTRON_SQUARED diff --git a/tardis/energy_input/util.py b/tardis/energy_input/util.py index 65e0400029b..d50f5893d35 100644 --- a/tardis/energy_input/util.py +++ b/tardis/energy_input/util.py @@ -4,6 +4,7 @@ from numba import njit from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel +from tardis.montecarlo.montecarlo_numba.opacities import kappa_calculation R_ELECTRON_SQUARED = (const.a0.cgs.value * const.alpha.cgs.value**2.0) ** 2.0 ELECTRON_MASS_ENERGY_KEV = (const.m_e * const.c**2.0).to("keV").value @@ -104,25 +105,6 @@ def angle_aberration_gamma(direction_vector, position_vector, time): return output_vector -@njit(**njit_dict_no_parallel) -def kappa_calculation(energy): - """ - Calculates kappa for various other calculations - i.e. energy normalized to electron rest energy - 511.0 KeV - - Parameters - ---------- - energy : float - - Returns - ------- - kappa : float - - """ - return energy / ELECTRON_MASS_ENERGY_KEV - - @njit(**njit_dict_no_parallel) def euler_rodrigues(theta, direction): """ diff --git a/tardis/montecarlo/montecarlo_numba/opacities.py b/tardis/montecarlo/montecarlo_numba/opacities.py index dd10595281f..70ab3f201ef 100644 --- a/tardis/montecarlo/montecarlo_numba/opacities.py +++ b/tardis/montecarlo/montecarlo_numba/opacities.py @@ -3,7 +3,6 @@ from numba import njit from tardis import constants as const -from tardis.energy_input.util import kappa_calculation from tardis.montecarlo.montecarlo_numba import ( njit_dict_no_parallel, ) @@ -21,12 +20,32 @@ MASS_FE = rd.Nuclide("Fe-56").atomic_mass * M_P SIGMA_T = const.sigma_T.cgs.value FINE_STRUCTURE = const.alpha.value +ELECTRON_MASS_ENERGY_KEV = (const.m_e * const.c**2.0).to("keV").value FF_OPAC_CONST = ( (2 * np.pi / (3 * M_E * K_B)) ** 0.5 * 4 * E**6 / (3 * M_E * H * C) ) # See Eq. 6.1.8 in http://personal.psu.edu/rbc3/A534/lec6.pdf +@njit(**njit_dict_no_parallel) +def kappa_calculation(energy): + """ + Calculates kappa for various other calculations + i.e. energy normalized to electron rest energy + 511.0 KeV + + Parameters + ---------- + energy : float + + Returns + ------- + kappa : float + + """ + return energy / ELECTRON_MASS_ENERGY_KEV + + @njit(**njit_dict_no_parallel) def chi_electron_calculator(opacity_state, nu, shell): """ diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_opacities.py b/tardis/montecarlo/montecarlo_numba/tests/test_opacities.py index dd9e65db7fa..d5143e296ce 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_opacities.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_opacities.py @@ -1,7 +1,12 @@ -import pytest import numpy.testing as npt +import pytest -import tardis.montecarlo.montecarlo_numba.opacities as calculate_opacity +from tardis.montecarlo.montecarlo_numba.opacities import ( + compton_opacity_calculation, + kappa_calculation, + pair_creation_opacity_calculation, + photoabsorption_opacity_calculation, +) @pytest.mark.parametrize( @@ -19,9 +24,7 @@ def test_compton_opacity_calculation(energy, electron_number_density, expected): energy : float electron_number_density : float """ - opacity = calculate_opacity.compton_opacity_calculation( - energy, electron_number_density - ) + opacity = compton_opacity_calculation(energy, electron_number_density) npt.assert_almost_equal(opacity, expected) @@ -45,7 +48,7 @@ def test_photoabsorption_opacity_calculation( ejecta_density : float iron_group_fraction : float """ - opacity = calculate_opacity.photoabsorption_opacity_calculation( + opacity = photoabsorption_opacity_calculation( energy, ejecta_density, iron_group_fraction ) @@ -71,8 +74,29 @@ def test_pair_creation_opacity_calculation( ejecta_density : float iron_group_fraction : float """ - opacity = calculate_opacity.pair_creation_opacity_calculation( + opacity = pair_creation_opacity_calculation( energy, ejecta_density, iron_group_fraction ) npt.assert_almost_equal(opacity, expected) + + +@pytest.mark.parametrize( + ["energy", "expected"], + [ + (511.0, 1.0000021334560507), + (255.5, 0.5000010667280254), + (0.0, 0.0), + (511.0e7, 10000021.334560508), + ], +) +def test_kappa_calculation(energy, expected): + """ + + Parameters + ---------- + energy : float + expected : float + """ + kappa = kappa_calculation(energy) + npt.assert_almost_equal(kappa, expected)