Formatting and another couple of fixes

wkerzendorf · Jul 12, 2024 · cf0512a · cf0512a
1 parent 611def4
commit cf0512a
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 15 deletions.
diff --git a/tardis/transport/montecarlo/base.py b/tardis/transport/montecarlo/base.py
@@ -187,9 +187,9 @@ def run(
             montecarlo_configuration=self.montecarlo_configuration,
         )
 
-        transport_state._montecarlo_virtual_luminosity.value[:] = (
-            v_packets_energy_hist
-        )
+        transport_state._montecarlo_virtual_luminosity.value[
+            :
+        ] = v_packets_energy_hist
         transport_state.last_interaction_type = last_interaction_tracker.types
         transport_state.last_interaction_in_nu = last_interaction_tracker.in_nus
         transport_state.last_line_interaction_in_id = (

diff --git a/tardis/transport/montecarlo/estimators/radfield_estimator_calcs.py b/tardis/transport/montecarlo/estimators/radfield_estimator_calcs.py
@@ -68,9 +68,9 @@ def update_bound_free_estimators(
         photo_ion_rate_estimator_increment = (
             comov_energy * distance * x_sect_bfs[i] / comov_nu
         )
-        estimator_state.photo_ion_estimator[current_continuum, shell_id] += (
-            photo_ion_rate_estimator_increment
-        )
+        estimator_state.photo_ion_estimator[
+            current_continuum, shell_id
+        ] += photo_ion_rate_estimator_increment
         estimator_state.stim_recomb_estimator[current_continuum, shell_id] += (
             photo_ion_rate_estimator_increment * boltzmann_factor
         )
@@ -82,12 +82,12 @@ def update_bound_free_estimators(
         bf_heating_estimator_increment = (
             comov_energy * distance * x_sect_bfs[i] * (1 - nu_th / comov_nu)
         )
-        estimator_state.bf_heating_estimator[current_continuum, shell_id] += (
-            bf_heating_estimator_increment
-        )
+        estimator_state.bf_heating_estimator[
+            current_continuum, shell_id
+        ] += bf_heating_estimator_increment
         estimator_state.stim_recomb_cooling_estimator[
             current_continuum, shell_id
-        ] += bf_heating_estimator_increment * boltzmann_factor
+        ] += (bf_heating_estimator_increment * boltzmann_factor)
 
 
 @njit(**njit_dict_no_parallel)
@@ -116,7 +116,7 @@ def update_line_estimators(
 
     radfield_mc_estimators.j_blue_estimator[
         cur_line_id, r_packet.current_shell_id
-    ] += energy / r_packet.nu
+    ] += (energy / r_packet.nu)
     radfield_mc_estimators.Edotlu_estimator[
         cur_line_id, r_packet.current_shell_id
     ] += energy
diff --git a/tardis/transport/montecarlo/packet_source.py b/tardis/transport/montecarlo/packet_source.py
@@ -119,7 +119,11 @@ def calculate_radfield_luminosity(self):
         astropy.units.Quantity
         """
         return (
-            4 * np.pi * const.sigma_sb * self.radius**2 * self.temperature**4
+            4
+            * np.pi
+            * const.sigma_sb
+            * self.radius**2
+            * self.temperature**4
         ).to("erg/s")
 
 
@@ -271,7 +275,8 @@ def set_temperature_from_luminosity(self, luminosity: u.Quantity):
 
         """
         self.temperature = (
-            (luminosity / (4 * np.pi * self.radius**2 * const.sigma_sb)) ** 0.25
+            (luminosity / (4 * np.pi * self.radius**2 * const.sigma_sb))
+            ** 0.25
         ).to("K")
 
 

diff --git a/tardis/transport/montecarlo/tests/test_interaction.py b/tardis/transport/montecarlo/tests/test_interaction.py
@@ -13,7 +13,7 @@ def test_thomson_scatter(packet, verysimple_time_explosion):
     init_energy = packet.energy
     time_explosion = verysimple_time_explosion
 
-    interaction.thomson_scatter(packet, time_explosion, False)
+    interaction.thomson_scatter(packet, time_explosion)
 
     assert np.abs(packet.mu - init_mu) > 1e-7
     assert np.abs(packet.nu - init_nu) > 1e-7

diff --git a/tardis/transport/montecarlo/tests/test_montecarlo_main_loop.py b/tardis/transport/montecarlo/tests/test_montecarlo_main_loop.py
@@ -6,6 +6,7 @@
 from tardis.simulation import Simulation
 from tardis.transport.montecarlo.configuration.montecarlo_globals import (
     LEGACY_MODE_ENABLED,
+    ENABLE_VPACKET_TRACKING,
 )
 
 
@@ -86,6 +87,7 @@ def test_montecarlo_main_loop_vpacket_log(
     regression_data,
     atomic_dataset,
 ):
+    ENABLE_VPACKET_TRACKING = True
     atomic_dataset = deepcopy(atomic_dataset)
     montecarlo_main_loop_config.montecarlo.no_of_virtual_packets = 5
 
@@ -99,7 +101,7 @@ def test_montecarlo_main_loop_vpacket_log(
 
     transport = montecarlo_main_loop_simulation.transport
 
-    assert transport.montecarlo_configuration.ENABLE_VPACKET_TRACKING is True
+    assert ENABLE_VPACKET_TRACKING is True
 
     expected_hdf_store = regression_data.sync_hdf_store(
         montecarlo_main_loop_simulation