Implement last interaction tracker

tardis/conftest.py

@@ -10,6 +10,7 @@
 from tardis.simulation import Simulation
 from tardis.tests.fixtures.atom_data import *
 from tardis.tests.fixtures.regression_data import regression_data
+from tardis import run_tardis
 
 # ensuring that regression_data is not removed by ruff
 assert regression_data is not None
@@ -219,6 +220,14 @@ def config_verysimple(example_configuration_dir):
     )
 
 
+# Creating a new config object so as to not overwrite other config.
+@pytest.fixture(scope="session")
+def config_rpacket_tracking(example_configuration_dir):
+    return Configuration.from_yaml(
+        example_configuration_dir / "tardis_configv1_verysimple.yml"
+    )
+
+
 @pytest.fixture(scope="function")
 def config_montecarlo_1e5_verysimple(example_configuration_dir):
     return Configuration.from_yaml(
@@ -243,3 +252,32 @@ def simulation_verysimple_vpacket_tracking(config_verysimple, atomic_dataset):
     sim.last_no_of_packets = 4000
     sim.run_final()
     return sim
+
+
+@pytest.fixture(scope="session")
+def simulation_rpacket_tracking(config_rpacket_tracking, atomic_dataset):
+    """
+    Creating a simulation object using a simple configuration
+
+    Parameters
+    ----------
+    config_verysimple_rpacket_tracking : tardis.io.configuration.config_reader.Configuration
+    atomic_dataset : AtomData
+
+    Returns
+    -------
+    simulation object with track_rpacket enabled
+    """
+    config_rpacket_tracking.montecarlo.iterations = 3
+    config_rpacket_tracking.montecarlo.no_of_packets = 4000
+    config_rpacket_tracking.montecarlo.last_no_of_packets = -1
+
+    config_rpacket_tracking.montecarlo.tracking.track_rpacket = True
+
+    atomic_data = deepcopy(atomic_dataset)
+    sim = run_tardis(
+        config_rpacket_tracking,
+        atom_data=atomic_data,
+        show_convergence_plots=False,
+    )
+    return sim

tardis/transport/montecarlo/__init__.py

@@ -22,9 +22,6 @@
 }
 
 from tardis.transport.montecarlo.r_packet import RPacket
-from tardis.transport.montecarlo.montecarlo_main_loop import (
-    montecarlo_main_loop,
-)
 from tardis.transport.montecarlo.packet_collections import (
     PacketCollection,
 )

tardis/transport/montecarlo/base.py

@@ -7,10 +7,9 @@
 from tardis import constants as const
 from tardis.io.logger import montecarlo_tracking as mc_tracker
 from tardis.io.util import HDFWriterMixin
-from tardis.transport.montecarlo import (
+from tardis.transport.montecarlo.montecarlo_main_loop import (
     montecarlo_main_loop,
 )
-from tardis.transport.montecarlo.configuration import montecarlo_globals
 from tardis.transport.montecarlo.configuration.base import (
     MonteCarloConfiguration,
     configuration_initialize,
@@ -210,11 +209,13 @@
 
         update_iterations_pbar(1)
         refresh_packet_pbar()
-        # Condition for Checking if RPacket Tracking is enabled
-        if self.enable_rpacket_tracking:
-            transport_state.rpacket_tracker = rpacket_trackers
 
-        if self.transport_state.rpacket_tracker is not None:
+        transport_state.rpacket_tracker = rpacket_trackers
+
+        # Need to change the implementation of rpacket_trackers_to_dataframe
+        # Such that it also takes of the case of
+        # RPacketLastInteractionTracker
+        if self.enable_rpacket_tracking:
             self.transport_state.rpacket_tracker_df = (
                 rpacket_trackers_to_dataframe(
                     self.transport_state.rpacket_tracker

tardis/transport/montecarlo/configuration/base.py

@@ -4,6 +4,7 @@
 import numpy as np
 
 from tardis.transport.montecarlo.configuration import montecarlo_globals
+from tardis.transport.montecarlo import montecarlo_main_loop
 from tardis.transport.montecarlo.numba_interface import (
     LineInteractionType,
 )
@@ -80,3 +81,6 @@
     montecarlo_globals.ENABLE_RPACKET_TRACKING = (
         transport.enable_rpacket_tracking
     )
+    montecarlo_main_loop.ENABLE_RPACKET_TRACKING = (
+        transport.enable_rpacket_tracking
+    )

tardis/transport/montecarlo/montecarlo_main_loop.py

@@ -10,7 +10,11 @@
     consolidate_vpacket_tracker,
     initialize_last_interaction_tracker,
 )
-from tardis.transport.montecarlo.packet_trackers import RPacketTracker
+import tardis.transport.montecarlo.montecarlo_main_loop as montecarlo_loop
+from tardis.transport.montecarlo.packet_trackers import (
+    RPacketTracker,
+    RPacketLastInteractionTracker,
+)
 from tardis.transport.montecarlo.r_packet import (
     PacketStatus,
     RPacket,
@@ -20,6 +24,8 @@
 )
 from tardis.util.base import update_packet_pbar
 
+ENABLE_RPACKET_TRACKING = False
+
 
 @njit(**njit_dict)
 def montecarlo_main_loop(
@@ -73,6 +79,17 @@
     vpacket_collections = List()
     # Configuring the Tracking for R_Packets
     rpacket_trackers = List()
+    if ENABLE_RPACKET_TRACKING:
+        for i in range(no_of_packets):
+            rpacket_trackers.append(
+                RPacketTracker(
+                    montecarlo_configuration.INITIAL_TRACKING_ARRAY_LENGTH
+                )
+            )
+    else:
+        for i in range(no_of_packets):
+            rpacket_trackers.append(RPacketLastInteractionTracker())
+
     for i in range(no_of_packets):
         vpacket_collections.append(
             VPacketCollection(
@@ -84,11 +101,6 @@
                 montecarlo_configuration.TEMPORARY_V_PACKET_BINS,
             )
         )
-        rpacket_trackers.append(
-            RPacketTracker(
-                montecarlo_configuration.INITIAL_TRACKING_ARRAY_LENGTH
-            )
-        )
 
     # Get the ID of the main thread and the number of threads
     main_thread_id = get_thread_id()
@@ -128,7 +140,6 @@
 
         # Get the local v_packet_collection for this thread
         vpacket_collection = vpacket_collections[i]
-
         # RPacket Tracker for this thread
         rpacket_tracker = rpacket_trackers[i]
 
@@ -186,7 +197,7 @@
             1,
         )
 
-    if montecarlo_globals.ENABLE_RPACKET_TRACKING:
+    if ENABLE_RPACKET_TRACKING:
         for rpacket_tracker in rpacket_trackers:
             rpacket_tracker.finalize_array()
 

tardis/transport/montecarlo/packet_trackers.py

@@ -35,7 +35,7 @@
         r : float
             Radius of the shell where the RPacket is present
         nu : float
-            Luminosity of the RPacket
+            Frequency of the RPacket
         mu : float
             Cosine of the angle made by the direction of movement of the RPacket from its original direction
         energy : float
@@ -156,3 +156,50 @@
         index=pd.MultiIndex.from_arrays(index_array, names=["index", "step"]),
         columns=df_dtypes.names,
     )
+
+
+rpacket_last_interaction_tracker_spec = [
+    ("index", int64),
+    ("r", float64),
+    ("nu", float64),
+    ("energy", float64),
+    ("shell_id", int64),
+    ("interaction_type", int64),
+]
+
+
+@jitclass(rpacket_last_interaction_tracker_spec)
+class RPacketLastInteractionTracker(object):
+    """
+    Numba JITCLASS for storing the last interaction the RPacket undergoes.
+    Parameters
+    ----------
+        index : int
+            Index position of each RPacket
+        r : float
+            Radius of the shell where the RPacket is present
+        nu : float
+            Frequency of the RPacket
+        energy : float
+            Energy possessed by the RPacket
+        shell_id : int
+            Current Shell No in which the last interaction happened
+        interaction_type: int
+            Type of interaction the rpacket undergoes
+    """
+
+    def __init__(self):
+        self.index = -1
+        self.r = -1.0
+        self.nu = 0.0
+        self.energy = 0.0
+        self.shell_id = -1
+        self.interaction_type = -1
+
+    def track(self, r_packet):
+        self.index = r_packet.index
+        self.r = r_packet.r
+        self.nu = r_packet.nu
+        self.energy = r_packet.energy
+        self.shell_id = r_packet.current_shell_id
+        self.interaction_type = r_packet.last_interaction_type

tardis/transport/montecarlo/single_packet_loop.py

@@ -9,6 +9,7 @@
     get_doppler_factor,
     get_inverse_doppler_factor,
 )
+from tardis.transport.montecarlo.r_packet import RPacket
 from tardis.transport.montecarlo.configuration import montecarlo_globals
 from tardis.transport.montecarlo.estimators.radfield_estimator_calcs import (
     update_bound_free_estimators,
@@ -83,8 +84,7 @@
         montecarlo_configuration.SURVIVAL_PROBABILITY,
     )
 
-    if montecarlo_globals.ENABLE_RPACKET_TRACKING:
-        rpacket_tracker.track(r_packet)
+    rpacket_tracker.track(r_packet)
 
     # this part of the code is temporary and will be better incorporated
     while r_packet.status == PacketStatus.IN_PROCESS:
@@ -257,9 +257,25 @@
             )
         else:
             pass
-        if montecarlo_globals.ENABLE_RPACKET_TRACKING:
+        if interaction_type != InteractionType.BOUNDARY:
             rpacket_tracker.track(r_packet)
 
+    # Registering the final boundary interaction.
+    # Only for RPacketTracker
+    # This is required by the RPacketPlotter tool
+    if montecarlo_globals.ENABLE_RPACKET_TRACKING:
+        temp_r_packet = RPacket(
+            r_packet.r,
+            r_packet.mu,
+            r_packet.nu,
+            r_packet.energy,
+            r_packet.seed,
+            r_packet.index,
+        )
+        temp_r_packet.current_shell_id = r_packet.current_shell_id
+        temp_r_packet.status = r_packet.status
+        rpacket_tracker.track(temp_r_packet)
+
 
 @njit
 def set_packet_props_partial_relativity(r_packet, time_explosion):