diff --git a/docs/io/output/rpacket_tracking.ipynb b/docs/io/output/rpacket_tracking.ipynb index ab61abe1913..168727da619 100644 --- a/docs/io/output/rpacket_tracking.ipynb +++ b/docs/io/output/rpacket_tracking.ipynb @@ -33,6 +33,7 @@ "
  • mu - Propagation Direction of the Packet (cosine of the angle the packet’s path makes with the radial direction)
    • \n", "
  • energy - Energy of the Packet
    • \n", "
  • shell_id - Current Shell Id where the Packet is present
    • \n", + "
  • interaction_type - Last Interaction type of the packet
    • \n", " \n", " \n", "" @@ -209,6 +210,7 @@ " mu\n", " energy\n", " shell_id\n", + " interaction_type\n", "```" ] }, @@ -298,7 +300,7 @@ "id": "ea308a55", "metadata": {}, "source": [ - "Thus, all other properties (`r`, `nu`, `mu`, `energy`, `shell_id`) can be accessed accordingly." + "Thus, all other properties (`r`, `nu`, `mu`, `energy`, `shell_id`,`interaction_type`) can be accessed accordingly." ] }, { @@ -402,6 +404,7 @@ "
  • mu
    • \n", "
  • energy
    • \n", "
  • shell_id
    • \n", + "
  • interaction_type
    • \n", " \n", " \n", "" @@ -503,7 +506,7 @@ "id": "6921f480", "metadata": {}, "source": [ - "Thus, all other properties (`status`, `seed`, `r`, `nu`, `mu`, `shell_id`) can be accessed accordingly." + "Thus, all other properties (`status`, `seed`, `r`, `nu`, `mu`, `shell_id`,`interaction_type`) can be accessed accordingly." ] } ], diff --git a/tardis/montecarlo/montecarlo_numba/numba_interface.py b/tardis/montecarlo/montecarlo_numba/numba_interface.py index 884991e45a1..455fb4419ce 100644 --- a/tardis/montecarlo/montecarlo_numba/numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/numba_interface.py @@ -431,6 +431,7 @@ def set_properties( ("mu", float64[:]), ("energy", float64[:]), ("shell_id", int64[:]), + ("interaction_type", int64[:]), ("interact_id", int64), ] @@ -460,6 +461,8 @@ class RPacketTracker(object): Energy possessed by the RPacket at a particular shell shell_id : int Current Shell No in which the RPacket is present + interaction_type: int + Type of interaction the rpacket undergoes interact_id : int Internal counter for the interactions that a particular RPacket undergoes """ @@ -474,6 +477,7 @@ def __init__(self): self.mu = np.empty(self.length, dtype=np.float64) self.energy = np.empty(self.length, dtype=np.float64) self.shell_id = np.empty(self.length, dtype=np.int64) + self.interaction_type = np.empty(self.length, dtype=np.int64) self.interact_id = 0 def track(self, r_packet): @@ -485,6 +489,7 @@ def track(self, r_packet): temp_mu = np.empty(temp_length, dtype=np.float64) temp_energy = np.empty(temp_length, dtype=np.float64) temp_shell_id = np.empty(temp_length, dtype=np.int64) + temp_interaction_type = np.empty(temp_length, dtype=np.int64) temp_status[: self.length] = self.status temp_r[: self.length] = self.r @@ -492,6 +497,7 @@ def track(self, r_packet): temp_mu[: self.length] = self.mu temp_energy[: self.length] = self.energy temp_shell_id[: self.length] = self.shell_id + temp_interaction_type[: self.length] = self.interaction_type self.status = temp_status self.r = temp_r @@ -499,6 +505,7 @@ def track(self, r_packet): self.mu = temp_mu self.energy = temp_energy self.shell_id = temp_shell_id + self.interaction_type = temp_interaction_type self.length = temp_length self.index = r_packet.index @@ -509,6 +516,7 @@ def track(self, r_packet): self.mu[self.interact_id] = r_packet.mu self.energy[self.interact_id] = r_packet.energy self.shell_id[self.interact_id] = r_packet.current_shell_id + self.interaction_type[self.interact_id] = r_packet.last_interaction_type self.interact_id += 1 def finalize_array(self): @@ -518,6 +526,7 @@ def finalize_array(self): self.mu = self.mu[: self.interact_id] self.energy = self.energy[: self.interact_id] self.shell_id = self.shell_id[: self.interact_id] + self.interaction_type = self.interaction_type[: self.interact_id] base_estimators_spec = [ diff --git a/tardis/montecarlo/montecarlo_numba/r_packet.py b/tardis/montecarlo/montecarlo_numba/r_packet.py index daa3a68e87b..88f985cf927 100644 --- a/tardis/montecarlo/montecarlo_numba/r_packet.py +++ b/tardis/montecarlo/montecarlo_numba/r_packet.py @@ -109,7 +109,7 @@ def rpacket_trackers_to_dataframe(rpacket_trackers): Returns ------- pandas.core.frame.DataFrame - Dataframe containing properties of RPackets as columns like status, seed, r, nu, mu, energy, shell_id + Dataframe containing properties of RPackets as columns like status, seed, r, nu, mu, energy, shell_id, interaction_type """ len_df = sum([len(tracker.r) for tracker in rpacket_trackers]) @@ -123,6 +123,7 @@ def rpacket_trackers_to_dataframe(rpacket_trackers): ("mu", np.float64), ("energy", np.float64), ("shell_id", np.int64), + ("interaction_type", np.int64), ] ) rpacket_tracker_ndarray = np.empty(len_df, df_dtypes) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_r_packet.py b/tardis/montecarlo/montecarlo_numba/tests/test_r_packet.py index 0ce23825001..d1959ac2747 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_r_packet.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_r_packet.py @@ -39,11 +39,22 @@ def test_rpacket_trackers_to_dataframe(simulation_rpacket_tracking_enabled): # check df shape and column names assert rtracker_df.shape == ( sum([len(tracker.r) for tracker in sim.runner.rpacket_tracker]), - 7, + 8, ) npt.assert_array_equal( sim.runner.rpacket_tracker_df.columns.values, - np.array(["status", "seed", "r", "nu", "mu", "energy", "shell_id"]), + np.array( + [ + "status", + "seed", + "r", + "nu", + "mu", + "energy", + "shell_id", + "interaction_type", + ] + ), ) # check all data with rpacket_tracker @@ -59,6 +70,7 @@ def test_rpacket_trackers_to_dataframe(simulation_rpacket_tracking_enabled): rpacket.mu[rpacket_step_no], rpacket.energy[rpacket_step_no], rpacket.shell_id[rpacket_step_no], + rpacket.interaction_type[rpacket_step_no], ] ) npt.assert_array_equal(rtracker_df.to_numpy(), np.array(expected_rtrackers))