Adding mu and r to virtual packet logging (tardis-sn#1696)

* added r and mu to vpacket logging

* fixing problem with merge conflict

* fixing test failure

* making documentation

docs/io/output/vpacket_logging.rst

@@ -25,6 +25,10 @@ After running the simulation, the following information can be retrieved:
       - List of virtual packet frequencies
     * - ``runner.virt_packet_energies``
       - List of virtual packet energies
+    * - ``runner.virt_packet_initial_mus``
+      - List of propagation directions that virtual packets are launched at
+    * - ``runner.virt_packet_initial_rs``
+      - List of radii that virtual packets are launched at
     * - ``runner.virt_packet_last_interaction_type``
       - Type of interaction that caused the virtual packet to be spawned
     * - ``runner.virt_packet_last_interaction_in_nu``

tardis/montecarlo/base.py

@@ -60,12 +60,14 @@ class MontecarloRunner(HDFWriterMixin):
     ]
 
     vpacket_hdf_properties = [
+        "virt_packet_nus",
+        "virt_packet_energies",
+        "virt_packet_initial_rs",
+        "virt_packet_initial_mus",
         "virt_packet_last_interaction_in_nu",
         "virt_packet_last_interaction_type",
         "virt_packet_last_line_interaction_in_id",
         "virt_packet_last_line_interaction_out_id",
-        "virt_packet_nus",
-        "virt_packet_energies",
     ]
 
     hdf_name = "runner"
@@ -129,6 +131,8 @@ def __init__(
         self.virt_packet_last_line_interaction_out_id = np.ones(2) * -1
         self.virt_packet_nus = np.ones(2) * -1.0
         self.virt_packet_energies = np.ones(2) * -1.0
+        self.virt_packet_initial_rs = np.ones(2) * -1.0
+        self.virt_packet_initial_mus = np.ones(2) * -1.0
 
         # set up logger based on config
         mc_logger.DEBUG_MODE = debug_packets

tardis/montecarlo/montecarlo_numba/base.py

@@ -62,6 +62,8 @@ def montecarlo_radial1d(model, plasma, runner):
         last_line_interaction_out_id,
         virt_packet_nus,
         virt_packet_energies,
+        virt_packet_initial_mus,
+        virt_packet_initial_rs,
         virt_packet_last_interaction_in_nu,
         virt_packet_last_interaction_type,
         virt_packet_last_line_interaction_in_id,
@@ -89,6 +91,12 @@ def montecarlo_radial1d(model, plasma, runner):
         runner.virt_packet_energies = np.concatenate(
             np.array(virt_packet_energies)
         ).ravel()
+        runner.virt_packet_initial_mus = np.concatenate(
+            np.array(virt_packet_initial_mus)
+        ).ravel()
+        runner.virt_packet_initial_rs = np.concatenate(
+            np.array(virt_packet_initial_rs)
+        ).ravel()
         runner.virt_packet_last_interaction_in_nu = np.concatenate(
             np.array(virt_packet_last_interaction_in_nu)
         ).ravel()
@@ -162,6 +170,8 @@ def montecarlo_main_loop(
     # Arrays for vpacket logging
     virt_packet_nus = []
     virt_packet_energies = []
+    virt_packet_initial_mus = []
+    virt_packet_initial_rs = []
     virt_packet_last_interaction_in_nu = []
     virt_packet_last_interaction_type = []
     virt_packet_last_line_interaction_in_id = []
@@ -204,6 +214,8 @@ def montecarlo_main_loop(
 
         vpackets_nu = vpacket_collection.nus[: vpacket_collection.idx]
         vpackets_energy = vpacket_collection.energies[: vpacket_collection.idx]
+        vpackets_initial_mu = vpacket_collection.initial_mus[: vpacket_collection.idx]
+        vpackets_initial_r = vpacket_collection.initial_rs[: vpacket_collection.idx]
 
         v_packets_idx = np.floor(
             (vpackets_nu - spectrum_frequency[0]) / delta_nu
@@ -221,17 +233,17 @@ def montecarlo_main_loop(
     if montecarlo_configuration.VPACKET_LOGGING:
         for vpacket_collection in vpacket_collections:
             vpackets_nu = vpacket_collection.nus[: vpacket_collection.idx]
-            vpackets_energy = vpacket_collection.energies[
-                : vpacket_collection.idx
-            ]
+            vpackets_energy = vpacket_collection.energies[: vpacket_collection.idx]
+            vpackets_initial_mu = vpacket_collection.initial_mus[: vpacket_collection.idx]
+            vpackets_initial_r = vpacket_collection.initial_rs[: vpacket_collection.idx]
             virt_packet_nus.append(np.ascontiguousarray(vpackets_nu))
             virt_packet_energies.append(np.ascontiguousarray(vpackets_energy))
-            virt_packet_last_interaction_in_nu.append(
-                np.ascontiguousarray(
-                    vpacket_collection.last_interaction_in_nu[
-                        : vpacket_collection.idx
-                    ]
-                )
+            virt_packet_initial_mus.append(np.ascontiguousarray(vpackets_initial_mu))
+            virt_packet_initial_rs.append(np.ascontiguousarray(vpackets_initial_r))
+            virt_packet_last_interaction_in_nu.append(np.ascontiguousarray(
+                vpacket_collection.last_interaction_in_nu[
+                    : vpacket_collection.idx
+                ])
             )
             virt_packet_last_interaction_type.append(
                 np.ascontiguousarray(
@@ -266,6 +278,8 @@ def montecarlo_main_loop(
         last_line_interaction_out_ids,
         virt_packet_nus,
         virt_packet_energies,
+        virt_packet_initial_mus,
+        virt_packet_initial_rs,
         virt_packet_last_interaction_in_nu,
         virt_packet_last_interaction_type,
         virt_packet_last_line_interaction_in_id,

tardis/montecarlo/montecarlo_numba/numba_interface.py

@@ -196,6 +196,8 @@ def __init__(
     ("v_packet_spawn_end_frequency", float64),
     ("nus", float64[:]),
     ("energies", float64[:]),
+    ("initial_mus", float64[:]),
+    ("initial_rs", float64[:]),
     ("idx", int64),
     ("number_of_vpackets", int64),
     ("length", int64),
@@ -222,6 +224,8 @@ def __init__(
         self.v_packet_spawn_end_frequency = v_packet_spawn_end_frequency
         self.nus = np.empty(temporary_v_packet_bins, dtype=np.float64)
         self.energies = np.empty(temporary_v_packet_bins, dtype=np.float64)
+        self.initial_mus = np.empty(temporary_v_packet_bins, dtype=np.float64)
+        self.initial_rs = np.empty(temporary_v_packet_bins, dtype=np.float64)
         self.number_of_vpackets = number_of_vpackets
         self.last_interaction_in_nu = np.zeros(
             temporary_v_packet_bins, dtype=np.float64
@@ -243,6 +247,8 @@ def set_properties(
         self,
         nu,
         energy,
+        initial_mu,
+        initial_r,
         last_interaction_in_nu,
         last_interaction_type,
         last_interaction_in_id,
@@ -252,6 +258,8 @@ def set_properties(
             temp_length = self.length * 2 + self.number_of_vpackets
             temp_nus = np.empty(temp_length, dtype=np.float64)
             temp_energies = np.empty(temp_length, dtype=np.float64)
+            temp_initial_mus = np.empty(temp_length, dtype=np.float64)
+            temp_initial_rs = np.empty(temp_length, dtype=np.float64)
             temp_last_interaction_in_nu = np.empty(
                 temp_length, dtype=np.float64
             )
@@ -260,6 +268,8 @@ def set_properties(
             temp_last_interaction_out_id = np.empty(temp_length, dtype=np.int64)
             temp_nus[: self.length] = self.nus
             temp_energies[: self.length] = self.energies
+            temp_initial_mus[: self.length] = self.initial_mus
+            temp_initial_rs[: self.length] = self.initial_rs
             temp_last_interaction_in_nu[
                 : self.length
             ] = self.last_interaction_in_nu
@@ -275,6 +285,8 @@ def set_properties(
 
             self.nus = temp_nus
             self.energies = temp_energies
+            self.initial_mus = temp_initial_mus
+            self.initial_rs = temp_initial_rs
             self.last_interaction_in_nu = temp_last_interaction_in_nu
             self.last_interaction_type = temp_last_interaction_type
             self.last_interaction_in_id = temp_last_interaction_in_id
@@ -283,6 +295,8 @@ def set_properties(
 
         self.nus[self.idx] = nu
         self.energies[self.idx] = energy
+        self.initial_mus[self.idx] = initial_mu
+        self.initial_rs[self.idx] = initial_r
         self.last_interaction_in_nu[self.idx] = last_interaction_in_nu
         self.last_interaction_type[self.idx] = last_interaction_type
         self.last_interaction_in_id[self.idx] = last_interaction_in_id

tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py

@@ -65,6 +65,8 @@ def test_VPacketCollection_set_properties(verysimple_3vpacket_collection):
 
     nus = [3.0e15, 0.0, 1e15, 1e5]
     energies = [0.4, 0.1, 0.6, 1e10]
+    initial_mus = [.1, 0, 1, .9]
+    initial_rs = [3e42, 4.5e45, 0, 9.0e40]
     last_interaction_in_nus = np.array(
         [3.0e15, 0.0, 1e15, 1e5], dtype=np.float64
     )
@@ -75,13 +77,17 @@ def test_VPacketCollection_set_properties(verysimple_3vpacket_collection):
     for (
         nu,
         energy,
+        initial_mu,
+        initial_r,
         last_interaction_in_nu,
         last_interaction_type,
         last_interaction_in_id,
         last_interaction_out_id,
     ) in zip(
         nus,
         energies,
+        initial_mus,
+        initial_rs,
         last_interaction_in_nus,
         last_interaction_types,
         last_interaction_in_ids,
@@ -90,6 +96,8 @@ def test_VPacketCollection_set_properties(verysimple_3vpacket_collection):
         verysimple_3vpacket_collection.set_properties(
             nu,
             energy,
+            initial_mu,
+            initial_r,
             last_interaction_in_nu,
             last_interaction_type,
             last_interaction_in_id,
@@ -108,6 +116,18 @@ def test_VPacketCollection_set_properties(verysimple_3vpacket_collection):
         ],
         energies,
     )
+    npt.assert_array_equal(
+        verysimple_3vpacket_collection.initial_mus[
+            : verysimple_3vpacket_collection.idx
+        ],
+        initial_mus,
+    )
+    npt.assert_array_equal(
+        verysimple_3vpacket_collection.initial_rs[
+            : verysimple_3vpacket_collection.idx
+        ],
+        initial_rs,
+    )
     npt.assert_array_equal(
         verysimple_3vpacket_collection.last_interaction_in_nu[
             : verysimple_3vpacket_collection.idx

tardis/montecarlo/montecarlo_numba/vpacket.py

@@ -273,6 +273,8 @@ def trace_vpacket_volley(
         vpacket_collection.set_properties(
             v_packet.nu,
             v_packet.energy,
+            v_packet_mu,
+            r_packet.r,
             r_packet.last_interaction_in_nu,
             r_packet.last_interaction_type,
             r_packet.last_line_interaction_in_id,