Add properties to transport state

tardis-sn · Apr 2, 2024 · d28fc89 · d28fc89
1 parent dc2cd8f
commit d28fc89
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 14 deletions.
diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_base.py b/tardis/montecarlo/montecarlo_numba/tests/test_base.py
@@ -50,13 +50,13 @@ def test_montecarlo_main_loop(
 
     # Load compare data from refdata
 
-    expected_nu = expected_hdf_store["/simulation/transport/output_nu"]
-    expected_energy = expected_hdf_store["/simulation/transport/output_energy"]
+    expected_nu = expected_hdf_store["/simulation/transport/transport_state/output_nu"]
+    expected_energy = expected_hdf_store["/simulation/transport/transport_state/output_energy"]
     expected_nu_bar_estimator = expected_hdf_store[
-        "/simulation/transport/nu_bar_estimator"
+        "/simulation/transport/transport_state/nu_bar_estimator"
     ]
     expected_j_estimator = expected_hdf_store[
-        "/simulation/transport/j_estimator"
+        "/simulation/transport/transport_state/j_estimator"
     ]
     expected_hdf_store.close()
     transport_state = montecarlo_main_loop_simulation.transport.transport_state
@@ -98,19 +98,19 @@ def test_montecarlo_main_loop_vpacket_log(
         montecarlo_main_loop_simulation
     )
 
-    expected_nu = expected_hdf_store["/simulation/transport/output_nu"]
-    expected_energy = expected_hdf_store["/simulation/transport/output_energy"]
+    expected_nu = expected_hdf_store["/simulation/transport/transport_state/output_nu"]
+    expected_energy = expected_hdf_store["/simulation/transport/transport_state/output_energy"]
     expected_nu_bar_estimator = expected_hdf_store[
-        "/simulation/transport/nu_bar_estimator"
+        "/simulation/transport/transport_state/nu_bar_estimator"
     ]
     expected_j_estimator = expected_hdf_store[
-        "/simulation/transport/j_estimator"
+        "/simulation/transport/transport_state/j_estimator"
     ]
     expected_vpacket_log_nus = expected_hdf_store[
-        "/simulation/transport/virt_packet_nus"
+        "/simulation/transport/transport_state/virt_packet_nus"
     ]
     expected_vpacket_log_energies = expected_hdf_store[
-        "/simulation/transport/virt_packet_energies"
+        "/simulation/transport/transport_state/virt_packet_energies"
     ]
 
     transport = montecarlo_main_loop_simulation.transport

diff --git a/tardis/montecarlo/montecarlo_transport_state.py b/tardis/montecarlo/montecarlo_transport_state.py
@@ -16,11 +16,13 @@ class MonteCarloTransportState(HDFWriterMixin):
         "output_energy",
         "nu_bar_estimator",
         "j_estimator",
+        "j_blue_estimator",
         "montecarlo_virtual_luminosity",
         "packet_luminosity",
         "spectrum",
         "spectrum_virtual",
         "spectrum_reabsorbed",
+        "spectrum_integrated",
         "time_of_simulation",
         "emitted_packet_mask",
         "last_interaction_type",
@@ -122,6 +124,10 @@ def nu_bar_estimator(self):
     def j_estimator(self):
         return self.radfield_mc_estimators.j_estimator
 
+    @property
+    def j_blue_estimator(self):
+        return self.radfield_mc_estimators.j_blue_estimator
+
     @property
     def time_of_simulation(self):
         return self.packet_collection.time_of_simulation * u.s

diff --git a/tardis/tests/test_tardis_full.py b/tardis/tests/test_tardis_full.py
@@ -72,7 +72,7 @@ def get_ref_data(key):
         return get_ref_data
 
     def test_j_blue_estimators(self, transport, refdata):
-        j_blue_estimator = refdata("j_blue_estimator").values
+        j_blue_estimator = refdata("transport_state/j_blue_estimator").values
 
         npt.assert_allclose(
             transport.transport_state.radfield_mc_estimators.j_blue_estimator,

diff --git a/tardis/tests/test_tardis_full_formal_integral.py b/tardis/tests/test_tardis_full_formal_integral.py
@@ -72,23 +72,23 @@ def get_ref_data(key):
         return get_ref_data
 
     def test_j_blue_estimators(self, transport, refdata):
-        j_blue_estimator = refdata("j_blue_estimator").values
+        j_blue_estimator = refdata("transport_state/j_blue_estimator").values
 
         npt.assert_allclose(
             transport.transport_state.radfield_mc_estimators.j_blue_estimator,
             j_blue_estimator,
         )
 
     def test_spectrum(self, transport, refdata):
-        luminosity = u.Quantity(refdata("spectrum/luminosity"), "erg /s")
+        luminosity = u.Quantity(refdata("transport_state/spectrum/luminosity"), "erg /s")
 
         assert_quantity_allclose(
             transport.transport_state.spectrum.luminosity, luminosity
         )
 
     def test_spectrum_integrated(self, transport, refdata):
         luminosity = u.Quantity(
-            refdata("spectrum_integrated/luminosity"), "erg /s"
+            refdata("transport_state/spectrum_integrated/luminosity"), "erg /s"
         )
 
         assert_quantity_allclose(