Add attributes specific to line interaction in the Tracker classes.

tardis/transport/montecarlo/base.py

@@ -173,6 +173,7 @@ def run(
             last_interaction_tracker,
             vpacket_tracker,
             rpacket_trackers,
+            rpacket_last_interaction_trackers,
         ) = montecarlo_main_loop(
             transport_state.packet_collection,
             transport_state.geometry_state,
@@ -212,13 +213,14 @@ def run(
         # Condition for Checking if RPacket Tracking is enabled
         if self.montecarlo_configuration.ENABLE_RPACKET_TRACKING:
             transport_state.rpacket_tracker = rpacket_trackers
-
-        if self.transport_state.rpacket_tracker is not None:
             self.transport_state.rpacket_tracker_df = (
                 rpacket_trackers_to_dataframe(
                     self.transport_state.rpacket_tracker
                 )
             )
+        else:
+            transport_state.rpacket_tracker = rpacket_last_interaction_trackers
+
         transport_state.virt_logging = (
             self.montecarlo_configuration.ENABLE_VPACKET_TRACKING
         )

tardis/transport/montecarlo/montecarlo_main_loop.py

@@ -4,7 +4,10 @@
 from numba.typed import List
 
 from tardis.transport.montecarlo import njit_dict
-from tardis.transport.montecarlo.packet_trackers import RPacketTracker
+from tardis.transport.montecarlo.packet_trackers import (
+    RPacketTracker,
+    RPacketLastInteractionTracker,
+)
 from tardis.transport.montecarlo.packet_collections import (
     VPacketCollection,
     consolidate_vpacket_tracker,
@@ -72,6 +75,7 @@ def montecarlo_main_loop(
     vpacket_collections = List()
     # Configuring the Tracking for R_Packets
     rpacket_trackers = List()
+    rpacket_last_interaction_trackers = List()
     for i in range(no_of_packets):
         vpacket_collections.append(
             VPacketCollection(
@@ -88,6 +92,9 @@ def montecarlo_main_loop(
                 montecarlo_configuration.INITIAL_TRACKING_ARRAY_LENGTH
             )
         )
+        rpacket_last_interaction_trackers.append(
+            RPacketLastInteractionTracker()
+        )
 
     # Get the ID of the main thread and the number of threads
     main_thread_id = get_thread_id()
@@ -130,6 +137,7 @@ def montecarlo_main_loop(
 
         # RPacket Tracker for this thread
         rpacket_tracker = rpacket_trackers[i]
+        rpacket_last_interaction_tracker = rpacket_last_interaction_trackers[i]
 
         loop = single_packet_loop(
             r_packet,
@@ -139,6 +147,7 @@ def montecarlo_main_loop(
             local_estimators,
             vpacket_collection,
             rpacket_tracker,
+            rpacket_last_interaction_tracker,
             montecarlo_configuration,
         )
         packet_collection.output_nus[i] = r_packet.nu
@@ -194,4 +203,5 @@ def montecarlo_main_loop(
         last_interaction_tracker,
         vpacket_tracker,
         rpacket_trackers,
+        rpacket_last_interaction_trackers,
     )

tardis/transport/montecarlo/packet_trackers.py

@@ -3,6 +3,8 @@
 import numpy as np
 import pandas as pd
 
+from tardis.transport.montecarlo.r_packet import InteractionType
+
 rpacket_tracker_spec = [
     ("length", int64),
     ("seed", int64),
@@ -14,6 +16,9 @@
     ("energy", float64[:]),
     ("shell_id", int64[:]),
     ("interaction_type", int64[:]),
+    ("interaction_in_line_nu", float64[:]),
+    ("interaction_in_line_id", int64[:]),
+    ("interaction_out_line_id", int64[:]),
     ("num_interactions", int64),
 ]
 
@@ -35,7 +40,7 @@ class RPacketTracker(object):
         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
@@ -44,6 +49,12 @@ class RPacketTracker(object):
             Current Shell No in which the RPacket is present
         interaction_type: int
             Type of interaction the rpacket undergoes
+        interaction_in_line_nu : float
+            frequency corresponding to the absroption line
+        interaction_in_line_id : int
+            Id of the absorption line
+        interaction_out_line_id : int
+            Id of the transmission line
         num_interactions : int
             Internal counter for the interactions that a particular RPacket undergoes
     """
@@ -59,6 +70,9 @@ def __init__(self, length):
         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.interaction_in_line_nu = np.empty(self.length, dtype=np.float64)
+        self.interaction_in_line_id = np.empty(self.length, dtype=np.int64)
+        self.interaction_out_line_id = np.empty(self.length, dtype=np.int64)
         self.num_interactions = 0
 
     def track(self, r_packet):
@@ -71,6 +85,11 @@ def track(self, r_packet):
             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_interaction_in_line_nu = np.empty(
+                temp_length, dtype=np.float64
+            )
+            temp_interaction_in_line_id = np.empty(temp_length, dtype=np.int64)
+            temp_interaction_out_line_id = np.empty(temp_length, dtype=np.int64)
 
             temp_status[: self.length] = self.status
             temp_r[: self.length] = self.r
@@ -79,6 +98,15 @@ def track(self, r_packet):
             temp_energy[: self.length] = self.energy
             temp_shell_id[: self.length] = self.shell_id
             temp_interaction_type[: self.length] = self.interaction_type
+            temp_interaction_in_line_nu[
+                : self.length
+            ] = self.interaction_in_line_nu
+            temp_interaction_in_line_id[
+                : self.length
+            ] = self.interaction_in_line_id
+            temp_interaction_out_line_id[
+                : self.length
+            ] = self.interaction_out_line_id
 
             self.status = temp_status
             self.r = temp_r
@@ -87,6 +115,9 @@ def track(self, r_packet):
             self.energy = temp_energy
             self.shell_id = temp_shell_id
             self.interaction_type = temp_interaction_type
+            self.interaction_in_line_nu = temp_interaction_in_line_nu
+            self.interaction_in_line_id = temp_interaction_in_line_id
+            self.interaction_out_line_id = temp_interaction_out_line_id
             self.length = temp_length
 
         self.index = r_packet.index
@@ -100,6 +131,21 @@ def track(self, r_packet):
         self.interaction_type[
             self.num_interactions
         ] = r_packet.last_interaction_type
+        # Only set if last interaction is line interaction, else -1 or 0.
+        if r_packet.last_interaction_type == InteractionType.LINE:
+            self.interaction_in_line_nu[
+                self.num_interactions
+            ] = r_packet.last_interaction_in_nu
+            self.interaction_in_line_id[
+                self.num_interactions
+            ] = r_packet.last_line_interaction_in_id
+            self.interaction_out_line_id[
+                self.num_interactions
+            ] = r_packet.last_line_interaction_out_id
+        else:
+            self.interaction_in_line_nu[self.num_interactions] = 0.0
+            self.interaction_in_line_id[self.num_interactions] = -1
+            self.interaction_out_line_id[self.num_interactions] = -1
         self.num_interactions += 1
 
     def finalize_array(self):
@@ -110,6 +156,15 @@ def finalize_array(self):
         self.energy = self.energy[: self.num_interactions]
         self.shell_id = self.shell_id[: self.num_interactions]
         self.interaction_type = self.interaction_type[: self.num_interactions]
+        self.interaction_in_line_nu = self.interaction_in_line_nu[
+            : self.num_interactions
+        ]
+        self.interaction_in_line_id = self.interaction_in_line_id[
+            : self.num_interactions
+        ]
+        self.interaction_out_line_id = self.interaction_out_line_id[
+            : self.num_interactions
+        ]
 
 
 def rpacket_trackers_to_dataframe(rpacket_trackers):
@@ -156,3 +211,70 @@ def rpacket_trackers_to_dataframe(rpacket_trackers):
         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),
+    ("interaction_in_line_nu", float64),
+    ("interaction_in_line_id", int64),
+    ("interaction_out_line_id", 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
+        interaction_in_line_nu : float
+            frequency corresponding to the absroption line
+        interaction_in_line_id : int
+            Id of the absorption line
+        interaction_out_line_id : int
+            Id of the transmission line
+    """
+
+    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
+        self.interaction_in_line_id = -1
+        self.interaction_out_line_id = -1
+        self.interaction_in_line_nu = 0.0
+
+    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
+        if r_packet.last_interaction_type == InteractionType.LINE:
+            self.interaction_in_line_nu = r_packet.last_interaction_in_nu
+            self.interaction_in_line_id = r_packet.last_line_interaction_in_id
+            self.interaction_out_line_id = r_packet.last_line_interaction_out_id
+        else:
+            self.interaction_in_line_nu = 0.0
+            self.interaction_in_line_id = -1
+            self.interaction_out_line_id = -1

tardis/transport/montecarlo/single_packet_loop.py

@@ -40,6 +40,7 @@ def single_packet_loop(
     estimators,
     vpacket_collection,
     rpacket_tracker,
+    rpacket_last_interaction_tracker,
     montecarlo_configuration,
 ):
     """
@@ -85,6 +86,8 @@ def single_packet_loop(
 
     if montecarlo_configuration.ENABLE_RPACKET_TRACKING:
         rpacket_tracker.track(r_packet)
+    else:
+        rpacket_last_interaction_tracker.track(r_packet)
 
     # this part of the code is temporary and will be better incorporated
     while r_packet.status == PacketStatus.IN_PROCESS:
@@ -264,8 +267,11 @@ def single_packet_loop(
             )
         else:
             pass
-        if montecarlo_configuration.ENABLE_RPACKET_TRACKING:
-            rpacket_tracker.track(r_packet)
+        if interaction_type != InteractionType.BOUNDARY:
+            if montecarlo_configuration.ENABLE_RPACKET_TRACKING:
+                rpacket_tracker.track(r_packet)
+            else:
+                rpacket_last_interaction_tracker.track(r_packet)
 
 
 @njit

tardis/transport/montecarlo/tests/test_r_packet.py

@@ -4,33 +4,42 @@
 import numpy.testing as npt
 import pytest
 
+from tardis.io.configuration.config_reader import Configuration
 from tardis.base import run_tardis
 from tardis.transport.montecarlo.packet_trackers import (
     rpacket_trackers_to_dataframe,
 )
 
 
-@pytest.fixture(scope="module")
-def simulation_rpacket_tracking_enabled(config_verysimple, atomic_dataset):
-    config_verysimple.montecarlo.iterations = 3
-    config_verysimple.montecarlo.no_of_packets = 4000
-    config_verysimple.montecarlo.last_no_of_packets = -1
-    config_verysimple.spectrum.virtual.virtual_packet_logging = True
-    config_verysimple.montecarlo.no_of_virtual_packets = 1
-    config_verysimple.montecarlo.tracking.track_rpacket = True
-    config_verysimple.spectrum.num = 2000
+@pytest.fixture()
+def config_rpacket_tracker(example_configuration_dir):
+    """Config object for rpacket tracker"""
+    return Configuration.from_yaml(
+        example_configuration_dir / "tardis_configv1_verysimple.yml"
+    )
+
+
+@pytest.fixture()
+def simulation_rpacket_tracking_enabled(config_rpacket_tracker, atomic_dataset):
+    """Simulation object with track_rpacket enabled"""
+    config_rpacket_tracker.montecarlo.iterations = 3
+    config_rpacket_tracker.montecarlo.no_of_packets = 4000
+    config_rpacket_tracker.montecarlo.last_no_of_packets = -1
+    config_rpacket_tracker.montecarlo.tracking.track_rpacket = True
+    config_rpacket_tracker.spectrum.num = 2000
     atomic_data = deepcopy(atomic_dataset)
     sim = run_tardis(
-        config_verysimple,
+        config_rpacket_tracker,
         atom_data=atomic_data,
         show_convergence_plots=False,
     )
     return sim
 
 
 def test_rpacket_trackers_to_dataframe(simulation_rpacket_tracking_enabled):
-    sim = simulation_rpacket_tracking_enabled
-    transport_state = sim.transport.transport_state
+    transport_state = (
+        simulation_rpacket_tracking_enabled.transport.transport_state
+    )
     rtracker_df = rpacket_trackers_to_dataframe(transport_state.rpacket_tracker)
 
     # check df shape and column names