Fix estimators parallel (#2030)

* Ran black * Co-authored-by: Sona Chitchyan <[email protected]> Co-authored-by: Ezequiel Pássaro <[email protected]> Co-authored-by: Laureano Martinez <[email protected]> * Further cleaned the code, removed commented out code, removed unused variables * Removed unused improts * Small formatting changes again * Removed extra ununsed variable * Removed unused estimator method * Ran black again
tardis-sn · May 19, 2022 · 0cd46f5 · 0cd46f5
1 parent f27fa30
commit 0cd46f5
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 9 deletions.
diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py
@@ -1,4 +1,9 @@
 from numba import prange, njit, objmode
+from numba.np.ufunc.parallel import (
+    _get_thread_id as get_thread_id,
+    get_num_threads,
+)
+
 import numpy as np
 
 from tardis.montecarlo.montecarlo_numba.r_packet import (
@@ -196,6 +201,24 @@ def montecarlo_main_loop(
         )
         rpacket_trackers.append(RPacketTracker())
 
+    main_thread_id = get_thread_id()
+    n_threads = get_num_threads()
+
+    estimator_list = List()
+    for i in range(n_threads):  # betting get tid goes from 0 to num threads
+        estimator_list.append(
+            Estimators(
+                np.copy(estimators.j_estimator),
+                np.copy(estimators.nu_bar_estimator),
+                np.copy(estimators.j_blue_estimator),
+                np.copy(estimators.Edotlu_estimator),
+                np.copy(estimators.photo_ion_estimator),
+                np.copy(estimators.stim_recomb_estimator),
+                np.copy(estimators.bf_heating_estimator),
+                np.copy(estimators.stim_recomb_cooling_estimator),
+                np.copy(estimators.photo_ion_estimator_statistics),
+            )
+        )
     # Arrays for vpacket logging
     virt_packet_nus = []
     virt_packet_energies = []
@@ -206,15 +229,18 @@ def montecarlo_main_loop(
     virt_packet_last_line_interaction_in_id = []
     virt_packet_last_line_interaction_out_id = []
     for i in prange(len(output_nus)):
+        tid = get_thread_id()
         if show_progress_bars:
-            with objmode:
-                update_amount = 1
-                update_packet_pbar(
-                    update_amount,
-                    current_iteration=iteration,
-                    no_of_packets=no_of_packets,
-                    total_iterations=total_iterations,
-                )
+
+            if tid == main_thread_id:
+                with objmode:
+                    update_amount = 1 * n_threads
+                    update_packet_pbar(
+                        update_amount,
+                        current_iteration=iteration,
+                        no_of_packets=no_of_packets,
+                        total_iterations=total_iterations,
+                    )
 
         if montecarlo_configuration.single_packet_seed != -1:
             seed = packet_seeds[montecarlo_configuration.single_packet_seed]
@@ -230,7 +256,7 @@ def montecarlo_main_loop(
             seed,
             i,
         )
-
+        local_estimators = estimator_list[tid]
         vpacket_collection = vpacket_collections[i]
         rpacket_tracker = rpacket_trackers[i]
 
@@ -277,6 +303,9 @@ def montecarlo_main_loop(
                 continue
             v_packets_energy_hist[idx] += vpackets_energy[j]
 
+    for sub_estimator in estimator_list:
+        estimators.increment(sub_estimator)
+
     if virtual_packet_logging:
         for vpacket_collection in vpacket_collections:
             vpackets_nu = vpacket_collection.nus[: vpacket_collection.idx]

diff --git a/tardis/montecarlo/montecarlo_numba/numba_interface.py b/tardis/montecarlo/montecarlo_numba/numba_interface.py
@@ -554,6 +554,22 @@ def __init__(
         self.stim_recomb_cooling_estimator = stim_recomb_cooling_estimator
         self.photo_ion_estimator_statistics = photo_ion_estimator_statistics
 
+    def increment(self, other):
+
+        self.j_estimator += other.j_estimator
+        self.nu_bar_estimator += other.nu_bar_estimator
+        self.j_blue_estimator += other.j_blue_estimator
+        self.Edotlu_estimator += other.Edotlu_estimator
+        self.photo_ion_estimator += other.photo_ion_estimator
+        self.stim_recomb_estimator += other.stim_recomb_estimator
+        self.bf_heating_estimator += other.bf_heating_estimator
+        self.stim_recomb_cooling_estimator += (
+            other.stim_recomb_cooling_estimator
+        )
+        self.photo_ion_estimator_statistics += (
+            other.photo_ion_estimator_statistics
+        )
+
 
 def configuration_initialize(runner, number_of_vpackets):
     if runner.line_interaction_type == "macroatom":