diff --git a/tardis/energy_input/gamma_ray_packet_source.py b/tardis/energy_input/gamma_ray_packet_source.py index f2a730dba2e..a5d1522043e 100644 --- a/tardis/energy_input/gamma_ray_packet_source.py +++ b/tardis/energy_input/gamma_ray_packet_source.py @@ -284,6 +284,7 @@ def create_packets(self, decays_per_isotope, *args, **kwargs): 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() @@ -721,29 +722,22 @@ def create_packets( 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(1)) + 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[ - sampled_packets_df["shell"].values - ] - sampled_packets_df["outer_velocity"] = self.outer_velocities[ - sampled_packets_df["shell"].values - ] + 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) - # packet decay time - times = self.create_packet_times_uniform_energy( - number_of_packets, - sampled_packets_df["isotopes"], - sampled_packets_df["time"], - ) - # get the time step index of the packets initial_time_indexes = sampled_packets_df.index.get_level_values(0) @@ -752,6 +746,13 @@ def create_packets( [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. @@ -764,9 +765,6 @@ def create_packets( # sample directions (valid at all times), non-relativistic directions = self.create_packet_directions(number_of_packets) - # get the packet shell index - shells = sampled_packets_df["shell"] - # the individual gamma-ray energy that makes up a packet # co-moving frame, including positronium formation nu_energies_cmf = self.create_packet_nus(