Refactoring adaptivity data handling to make it more cleaner and help…

… with logging

micro_manager/adaptivity/adaptivity.py

@@ -35,6 +35,9 @@ def __init__(self, configurator, logger) -> None:
             configurator.get_adaptivity_similarity_measure()
         )
 
+        self._active_sim_ids = np.array([], dtype=int)
+        self._inactive_sim_ids = np.array([], dtype=int)
+
     def _get_similarity_dists(
         self, dt: float, similarity_dists: np.ndarray, data: dict
     ) -> np.ndarray:

micro_manager/adaptivity/global_adaptivity.py

@@ -70,9 +70,7 @@ def compute_adaptivity(
         self,
         dt: float,
         micro_sims: list,
-        similarity_dists_nm1: np.ndarray,
-        is_sim_active_nm1: np.ndarray,
-        sim_is_associated_to_nm1: np.ndarray,
+        adaptivity_data_nm1: list,
         data_for_adaptivity: dict,
     ) -> tuple:
         """
@@ -84,21 +82,15 @@ def compute_adaptivity(
             Current time step of the macro-micro coupled problem
         micro_sims : list
             List of objects of class MicroProblem, which are the micro simulations
-        similarity_dists_nm1 : numpy array
-            2D array having similarity distances between each micro simulation pair
-        is_sim_active_nm1 : numpy array
-            1D array having state (active or inactive) of each micro simulation on this rank
-        sim_is_associated_to_nm1 : numpy array
-            1D array with values of associated simulations of inactive simulations. Active simulations have None
+        adaptivity_data_nm1 : list
+            List of numpy arrays: similarity_dists (2D array having similarity distances between each micro simulation pair), is_sim_active (1D array having state (active or inactive) of each micro simulation), sim_is_associated_to (1D array with values of associated simulations of inactive simulations. Active simulations have None)
         data_for_adaptivity : dict
             Dictionary with keys as names of data to be used in the similarity calculation, and values as the respective data for the micro simulations
 
         Results
         -------
-        similarity_dists : numpy array
-            2D array having similarity distances between each micro simulation pair
-        is_sim_active : numpy array
-            1D array having state (active or inactive) of each micro simulation
+        list
+            List of numpy arrays: similarity_dists (2D array having similarity distances between each micro simulation pair), is_sim_active (1D array having state (active or inactive) of each micro simulation), sim_is_associated_to (1D array with values of associated simulations of inactive simulations. Active simulations have None)
         """
         for name in data_for_adaptivity.keys():
             if name not in self._adaptivity_data_names:
@@ -115,18 +107,27 @@ def compute_adaptivity(
             global_data_for_adaptivity[name] = np.concatenate((data_as_list[:]), axis=0)
 
         similarity_dists = self._get_similarity_dists(
-            dt, similarity_dists_nm1, global_data_for_adaptivity
+            dt, adaptivity_data_nm1[0], global_data_for_adaptivity
         )
 
-        is_sim_active = self._update_active_sims(similarity_dists, is_sim_active_nm1)
+        is_sim_active = self._update_active_sims(
+            similarity_dists, adaptivity_data_nm1[1]
+        )
 
         is_sim_active, sim_is_associated_to = self._update_inactive_sims(
-            similarity_dists, is_sim_active, sim_is_associated_to_nm1, micro_sims
+            similarity_dists, is_sim_active, adaptivity_data_nm1[2], micro_sims
         )
         sim_is_associated_to = self._associate_inactive_to_active(
             similarity_dists, is_sim_active, sim_is_associated_to
         )
 
+        self._active_sim_ids = np.where(
+            is_sim_active[self._global_ids[0] : self._global_ids[-1] + 1]
+        )[0]
+        self._inactive_sim_ids = np.where(
+            is_sim_active[self._global_ids[0] : self._global_ids[-1] + 1] == False
+        )[0]
+
         self._logger.info(
             "{} active simulations, {} inactive simulations".format(
                 np.count_nonzero(
@@ -139,12 +140,11 @@ def compute_adaptivity(
             )
         )
 
-        return similarity_dists, is_sim_active, sim_is_associated_to
+        return [similarity_dists, is_sim_active, sim_is_associated_to]
 
     def communicate_micro_output(
         self,
-        is_sim_active: np.ndarray,
-        sim_is_associated_to: np.ndarray,
+        adaptivity_data: list,
         micro_output: list,
     ) -> None:
         """
@@ -155,13 +155,16 @@ def communicate_micro_output(
         ----------
         micro_sims : list
             List of objects of class MicroProblem, which are the micro simulations
-        is_sim_active : numpy array
-            1D array having state (active or inactive) of each micro simulation on this rank
-        sim_is_associated_to : numpy array
-            1D array with values of associated simulations of inactive simulations. Active simulations have -2
+        adaptivity_data : list
+            List of numpy arrays:
+                is_sim_active (1D array having state (active or inactive) of each micro simulation)
+                sim_is_associated_to (1D array with values of associated simulations of inactive simulations. Active simulations have None)
         micro_output : list
             List of dicts having individual output of each simulation. Only the active simulation outputs are entered.
         """
+        is_sim_active = adaptivity_data[0]
+        sim_is_associated_to = adaptivity_data[1]
+
         inactive_local_ids = np.where(
             is_sim_active[self._global_ids[0] : self._global_ids[-1] + 1] == False
         )[0]
@@ -198,6 +201,28 @@ def communicate_micro_output(
             for local_id in active_to_inactive_map[assoc_active_ids[count]]:
                 micro_output[local_id] = deepcopy(output)
 
+    def get_active_sim_ids(self) -> np.ndarray:
+        """
+        Get the ids of active simulations.
+
+        Returns
+        -------
+        numpy array
+            1D array of active simulation ids
+        """
+        return self._active_sim_ids
+
+    def get_inactive_sim_ids(self) -> np.ndarray:
+        """
+        Get the ids of inactive simulations.
+
+        Returns
+        -------
+        numpy array
+            1D array of inactive simulation ids
+        """
+        return self._inactive_sim_ids
+
     def _update_inactive_sims(
         self,
         similarity_dists: np.ndarray,

micro_manager/adaptivity/local_adaptivity.py

@@ -26,9 +26,7 @@ def compute_adaptivity(
         self,
         dt,
         micro_sims,
-        similarity_dists_nm1: np.ndarray,
-        is_sim_active_nm1: np.ndarray,
-        sim_is_associated_to_nm1: np.ndarray,
+        adaptivity_data_nm1: list,
         data_for_adaptivity: dict,
     ) -> tuple:
         """
@@ -40,22 +38,16 @@ def compute_adaptivity(
             Current time step
         micro_sims : list
             List containing simulation objects
-        similarity_dists_nm1 : numpy array
-            2D array having similarity distances between each micro simulation pair.
-        is_sim_active_nm1 : numpy array
-            1D array having True if sim is active, False if sim is inactive.
-        sim_is_associated_to_nm1 : numpy array
-            1D array with values of associated simulations of inactive simulations. Active simulations have None.
+        adaptivity_data_nm1 : list
+            List of numpy arrays: similarity_dists (2D array having similarity distances between each micro simulation pair), is_sim_active (1D array having state (active or inactive) of each micro simulation), sim_is_associated_to (1D array with values of associated simulations of inactive simulations. Active simulations have None)
         data_for_adaptivity : dict
             A dictionary containing the names of the data to be used in adaptivity as keys and information on whether
             the data are scalar or vector as values.
 
         Returns
         -------
-        similarity_dists : numpy array
-            2D array having similarity distances between each micro simulation pair.
-        is_sim_active : numpy array
-            1D array, True is sim is active, False if sim is inactive.
+        list
+            List of numpy arrays: similarity_dists (2D array having similarity distances between each micro simulation pair), is_sim_active (1D array having state (active or inactive) of each micro simulation), sim_is_associated_to (1D array with values of associated simulations of inactive simulations. Active simulations have None)
         """
         for name in data_for_adaptivity.keys():
             if name not in self._adaptivity_data_names:
@@ -66,14 +58,16 @@ def compute_adaptivity(
                 )
 
         similarity_dists = self._get_similarity_dists(
-            dt, similarity_dists_nm1, data_for_adaptivity
+            dt, adaptivity_data_nm1[0], data_for_adaptivity
         )
 
         # Operation done globally if global adaptivity is chosen
-        is_sim_active = self._update_active_sims(similarity_dists, is_sim_active_nm1)
+        is_sim_active = self._update_active_sims(
+            similarity_dists, adaptivity_data_nm1[1]
+        )
 
         is_sim_active, sim_is_associated_to = self._update_inactive_sims(
-            similarity_dists, is_sim_active, sim_is_associated_to_nm1, micro_sims
+            similarity_dists, is_sim_active, adaptivity_data_nm1[2], micro_sims
         )
 
         sim_is_associated_to = self._associate_inactive_to_active(
@@ -87,7 +81,32 @@ def compute_adaptivity(
             )
         )
 
-        return similarity_dists, is_sim_active, sim_is_associated_to
+        self._active_sim_ids = np.where(is_sim_active)[0]
+        self._inactive_sim_ids = np.where(is_sim_active == False)[0]
+
+        return [similarity_dists, is_sim_active, sim_is_associated_to]
+
+    def get_active_sim_ids(self) -> np.ndarray:
+        """
+        Get the ids of active simulations.
+
+        Returns
+        -------
+        numpy array
+            1D array of active simulation ids
+        """
+        return self._active_sim_ids
+
+    def get_inactive_sim_ids(self) -> np.ndarray:
+        """
+        Get the ids of inactive simulations.
+
+        Returns
+        -------
+        numpy array
+            1D array of inactive simulation ids
+        """
+        return self._inactive_sim_ids
 
     def _update_inactive_sims(
         self,