Make global adaptivity work for cases where ranks have unequal number of micro simulations

micro_manager/adaptivity/global_adaptivity.py

@@ -18,8 +18,7 @@ def __init__(
             self,
             configurator,
             logger,
-            is_sim_on_this_rank: list,
-            rank_of_sim: np.ndarray,
+            global_number_of_sims: float,
             global_ids: list,
             rank: int,
             comm) -> None:
@@ -32,10 +31,8 @@ def __init__(
             Object which has getter functions to get parameters defined in the configuration file.
         logger : object of logging
             Logger defined from the standard package logging
-        is_sim_on_this_rank : list
+        global_number_of_sims : float
             List of booleans. True if simulation is on this rank, False otherwise.
-        rank_of_sim : numpy array
-            1D array consisting of rank on which the simulation lives.
         global_ids : list
             List of global IDs of simulations living on this rank.
         rank : int
@@ -44,12 +41,26 @@ def __init__(
             Global communicator of MPI.
         """
         super().__init__(configurator, logger)
-        self._is_sim_on_this_rank = is_sim_on_this_rank
-        self._rank_of_sim = rank_of_sim
         self._global_ids = global_ids
         self._comm = comm
         self._rank = rank
 
+        local_number_of_sims = len(global_ids)
+
+        # Create a map of micro simulation global IDs and the ranks on which they are
+        micro_sims_on_this_rank = np.zeros(local_number_of_sims, dtype=np.intc)
+        for i in range(local_number_of_sims):
+            micro_sims_on_this_rank[i] = self._rank
+
+        self._rank_of_sim = np.zeros(global_number_of_sims, dtype=np.intc)  # DECLARATION
+
+        self._comm.Allgatherv(micro_sims_on_this_rank, self._rank_of_sim)
+
+        self._is_sim_on_this_rank = [False] * global_number_of_sims  # DECLARATION
+        for i in range(global_number_of_sims):
+            if self._rank_of_sim[i] == self._rank:
+                self._is_sim_on_this_rank[i] = True
+
     def compute_adaptivity(
             self,
             dt: float,

micro_manager/micro_manager.py

@@ -314,25 +314,15 @@ def _initialize(self) -> None:
                 fromlist=["MicroSimulation"]),
             "MicroSimulation")
 
-        if self._is_adaptivity_on:
-            # Create micro simulation objects
-            for i in range(self._local_number_of_sims):
-                self._micro_sims[i] = create_simulation_class(
-                    micro_problem)(self._global_ids_of_local_sims[i])
-
-            # Create a map of micro simulation global IDs and the ranks on which they are
-            micro_sims_on_this_rank = np.zeros(self._local_number_of_sims, dtype=np.intc)
-            for i in range(self._local_number_of_sims):
-                micro_sims_on_this_rank[i] = self._rank
-
-            self._rank_of_sim = np.zeros(self._global_number_of_sims, dtype=np.intc)  # DECLARATION
-            self._comm.Allgather(micro_sims_on_this_rank, self._rank_of_sim)
+        # Create micro simulation objects
+        for i in range(self._local_number_of_sims):
+            self._micro_sims[i] = create_simulation_class(
+                micro_problem)(self._global_ids_of_local_sims[i])
 
-            self._is_sim_on_this_rank = [False] * self._global_number_of_sims  # DECLARATION
-            for i in range(self._global_number_of_sims):
-                if self._rank_of_sim[i] == self._rank:
-                    self._is_sim_on_this_rank[i] = True
+        self._logger.info("Micro simulations with global IDs {} - {} created.".format(
+            self._global_ids_of_local_sims[0], self._global_ids_of_local_sims[-1]))
 
+        if self._is_adaptivity_on:
             if self._adaptivity_type == "local":
                 self._adaptivity_controller = LocalAdaptivityCalculator(
                     self._config, self._logger)
@@ -341,21 +331,13 @@ def _initialize(self) -> None:
                 self._adaptivity_controller = GlobalAdaptivityCalculator(
                     self._config,
                     self._logger,
-                    self._is_sim_on_this_rank,
-                    self._rank_of_sim,
+                    self._global_number_of_sims,
                     self._global_ids_of_local_sims,
                     self._rank,
                     self._comm)
                 self._number_of_sims_for_adaptivity = self._global_number_of_sims
 
             self._micro_sims_active_steps = np.zeros(self._local_number_of_sims)
-        else:
-            for i in range(self._local_number_of_sims):
-                self._micro_sims[i] = (
-                    create_simulation_class(micro_problem)(self._global_ids_of_local_sims[i]))
-
-        self._logger.info("Micro simulations with global IDs {} - {} created.".format(
-            self._global_ids_of_local_sims[0], self._global_ids_of_local_sims[-1]))
 
         self._micro_sims_have_output = False
         if hasattr(micro_problem, 'output') and callable(getattr(micro_problem, 'output')):

tests/unit/test_adaptivity_parallel.py

@@ -17,14 +17,11 @@ def test_update_inactive_sims_global_adaptivity(self):
         Run this test in parallel using MPI with 2 ranks.
         """
         if self._rank == 0:
-            is_sim_on_this_rank = [True, True, True, False, False]
             global_ids = [0, 1, 2]
         elif self._rank == 1:
-            is_sim_on_this_rank = [False, False, False, True, True]
             global_ids = [3, 4]
 
         is_sim_active = np.array([False, False, True, True, False])
-        rank_of_sim = [0, 0, 0, 1, 1]
         sim_is_associated_to = [3, 3, -2, -2, 2]
         expected_is_sim_active = np.array([True, False, True, True, True])
         expected_sim_is_associated_to = [-2, 3, -2, -2, -2]
@@ -34,8 +31,7 @@ def test_update_inactive_sims_global_adaptivity(self):
         adaptivity_controller = GlobalAdaptivityCalculator(
             configurator,
             MagicMock(),
-            is_sim_on_this_rank,
-            rank_of_sim,
+            5,
             global_ids,
             rank=self._rank,
             comm=self._comm)
@@ -87,27 +83,23 @@ def test_communicate_micro_output(self):
         output_1 = {"data1.1": 10.0, "data1.2": [10.0, 20.0]}
 
         if self._rank == 0:
-            is_sim_on_this_rank = [True, True, True, False, False]
             global_ids = [0, 1, 2]
             sim_output = [None, None, output_0]
             expected_sim_output = [output_1, output_1, output_0]
         elif self._rank == 1:
-            is_sim_on_this_rank = [False, False, False, True, True]
             global_ids = [3, 4]
             sim_output = [output_1, None]
             expected_sim_output = [output_1, output_0]
 
         is_sim_active = np.array([False, False, True, True, False])
-        rank_of_sim = [0, 0, 0, 1, 1]
         sim_is_associated_to = [3, 3, -2, -2, 2]
 
         configurator = MagicMock()
         configurator.get_adaptivity_similarity_measure = MagicMock(return_value="L1")
         adaptivity_controller = GlobalAdaptivityCalculator(
             configurator,
             MagicMock(),
-            is_sim_on_this_rank,
-            rank_of_sim,
+            5,
             global_ids,
             rank=self._rank,
             comm=self._comm)