Fix get_ranks_of_sims test and add load balancing threshold as config…

… parameter

micro_manager/adaptivity/global_adaptivity_lb.py

@@ -55,6 +55,8 @@ def __init__(
             configurator.is_load_balancing_two_step()
         )
 
+        self._threshold = configurator.get_load_balancing_threshold()
+
     def redistribute_sims(self, micro_sims: list) -> None:
         """
         Redistribute simulations among ranks to balance compute load.
@@ -87,8 +89,8 @@ def _redistribute_active_sims(self, micro_sims: list) -> None:
         psend_sims = 0
         precv_sims = 0
 
-        f_avg_active_sims = math.floor(avg_active_sims)
-        c_avg_active_sims = math.ceil(avg_active_sims)
+        f_avg_active_sims = math.floor(avg_active_sims) - self._threshold
+        c_avg_active_sims = math.ceil(avg_active_sims) + self._threshold
 
         if n_active_sims_local == f_avg_active_sims:
             # Simulations to potentially receive

micro_manager/config.py

@@ -56,6 +56,7 @@ def __init__(self, config_file_name):
         self._adaptivity_is_load_balancing = False
         self._load_balancing_n = 1
         self._two_step_load_balancing = False
+        self._load_balancing_threshold = 1
 
         # Snapshot information
         self._parameter_file_name = None
@@ -264,7 +265,7 @@ def read_json_micro_manager(self):
                 try:
                     self._load_balancing_n = self._data["simulation_params"][
                         "adaptivity_settings"
-                    ]["load_balancing_n"]
+                    ]["load_balancing_settings"]["load_balancing_n"]
                 except BaseException:
                     self._logger.log_info_one_rank(
                         "No load balancing interval provided. Load balancing will be performed every time window. THIS IS NOT RECOMMENDED."
@@ -273,8 +274,8 @@ def read_json_micro_manager(self):
                 try:
                     if (
                         self._data["simulation_params"]["adaptivity_settings"][
-                            "two_step_load_balancing"
-                        ]
+                            "load_balancing_settings"
+                        ]["two_step_load_balancing"]
                         == "True"
                     ):
                         self._two_step_load_balancing = True
@@ -283,6 +284,15 @@ def read_json_micro_manager(self):
                         "Two-step load balancing is not specified. Micro Manager will only try to balance the load in one sweep."  # TODO: Need a better log message here.
                     )
 
+                try:
+                    self._load_balancing_threshold = self._data["simulation_params"][
+                        "adaptivity_settings"
+                    ]["load_balancing_settings"]["load_balancing_threshold"]
+                except BaseException:
+                    self._logger.log_info_one_rank(
+                        "No load balancing threshold provided. The default threshold of 1 will be used."
+                    )
+
             self._write_data_names.append("active_state")
             self._write_data_names.append("active_steps")
 
@@ -631,6 +641,17 @@ def is_load_balancing_two_step(self):
         """
         return self._two_step_load_balancing
 
+    def get_load_balancing_threshold(self):
+        """
+        Get the load balancing threshold to control how balanced the micro simulations need to be.
+
+        Returns
+        -------
+        load_balancing_threshold : float
+            Load balancing threshold
+        """
+        return self._load_balancing_threshold
+
     def get_micro_dt(self):
         """
         Get the size of the micro time window.

tests/unit/test_adaptivity_parallel.py

@@ -13,6 +13,12 @@ def setUp(self):
         self._rank = self._comm.Get_rank()
         self._size = self._comm.Get_size()
 
+        self._configurator = MagicMock()
+        self._configurator.get_adaptivity_similarity_measure = MagicMock(
+            return_value="L1"
+        )
+        self._configurator.get_output_dir = MagicMock(return_value="output_dir")
+
     def test_update_inactive_sims_global_adaptivity(self):
         """
         Test functionality to update inactive simulations in a particular setting.
@@ -28,12 +34,8 @@ def test_update_inactive_sims_global_adaptivity(self):
         expected_is_sim_active = np.array([True, False, True, True, True])
         expected_sim_is_associated_to = [-2, 3, -2, -2, -2]
 
-        configurator = MagicMock()
-        configurator.get_adaptivity_similarity_measure = MagicMock(return_value="L1")
-        configurator.get_output_dir = MagicMock(return_value="output_dir")
-
         adaptivity_controller = GlobalAdaptivityCalculator(
-            configurator, 5, global_ids, rank=self._rank, comm=self._comm
+            self._configurator, 5, global_ids, rank=self._rank, comm=self._comm
         )
 
         # Force the activation of sim #0 and #4
@@ -105,15 +107,15 @@ def test_update_all_active_sims_global_adaptivity(self):
         expected_is_sim_active = np.array([False, False, False, False, True])
         expected_sim_is_associated_to = [4, 4, 4, 4, -2]
 
-        configurator = MagicMock()
-        configurator.get_adaptivity_hist_param = MagicMock(return_value=0.1)
-        configurator.get_adaptivity_refining_const = MagicMock(return_value=0.05)
-        configurator.get_adaptivity_coarsening_const = MagicMock(return_value=0.2)
-        configurator.get_adaptivity_similarity_measure = MagicMock(return_value="L2rel")
-        configurator.get_output_dir = MagicMock(return_value="output_dir")
+        self._configurator.get_adaptivity_hist_param = MagicMock(return_value=0.1)
+        self._configurator.get_adaptivity_refining_const = MagicMock(return_value=0.05)
+        self._configurator.get_adaptivity_coarsening_const = MagicMock(return_value=0.2)
+        self._configurator.get_adaptivity_similarity_measure = MagicMock(
+            return_value="L2rel"
+        )
 
         adaptivity_controller = GlobalAdaptivityCalculator(
-            configurator, 5, global_ids, rank=self._rank, comm=self._comm
+            self._configurator, 5, global_ids, rank=self._rank, comm=self._comm
         )
 
         adaptivity_controller._adaptivity_data_names = {
@@ -176,12 +178,8 @@ def test_communicate_micro_output(self):
         is_sim_active = np.array([False, False, True, True, False])  # is_sim_active
         sim_is_associated_to = [3, 3, -2, -2, 2]  # sim_is_associated_to
 
-        configurator = MagicMock()
-        configurator.get_adaptivity_similarity_measure = MagicMock(return_value="L1")
-        configurator.get_output_dir = MagicMock(return_value="output_dir")
-
         adaptivity_controller = GlobalAdaptivityCalculator(
-            configurator, 5, global_ids, rank=self._rank, comm=self._comm
+            self._configurator, 5, global_ids, rank=self._rank, comm=self._comm
         )
 
         adaptivity_controller._communicate_micro_output(
@@ -191,7 +189,9 @@ def test_communicate_micro_output(self):
         self.assertTrue(np.array_equal(expected_sim_output, sim_output))
 
     def test_get_ranks_of_sims(self):
-        """ """
+        """
+        ...
+        """
         if self._rank == 0:
             global_ids = [0, 1, 2]
             expected_ranks_of_sims = [0, 0, 0, 1, 1]

tests/unit/test_global_adaptivity_lb.py

@@ -374,8 +374,6 @@ def test_redistribute_inactive_sims_four_ranks(self):
             if micro_sims[i] is not None:
                 actual_global_ids.append(micro_sims[i].get_global_id())
 
-        # print("Rank {}: Actual global IDs: {}".format(self._rank, actual_global_ids))
-
         self.assertEqual(actual_global_ids, expected_global_ids)
 
         actual_ranks_of_sims = adaptivity_controller._get_ranks_of_sims()