Repository: DARMA-tasking/LB-analysis-framework

Issue:
Add a target imbalance threshold to yaml
Add an optional target imbalance threshold to the yaml file that can be used to conclude the iteration process in fewer iterations than requested iff the imbalance has been brought down below the threshold already.

Discussion:

Patch:
diff --git a/src/lbaf/Execution/lbsInformAndTransferAlgorithm.py b/src/lbaf/Execution/lbsInformAndTransferAlgorithm.py
index c854aca..4b21239 100644
--- a/src/lbaf/Execution/lbsInformAndTransferAlgorithm.py
+++ b/src/lbaf/Execution/lbsInformAndTransferAlgorithm.py
@@ -7,7 +7,6 @@ from .lbsCriterionBase import CriterionBase
 from .lbsTransferStrategyBase import TransferStrategyBase
 from ..Model.lbsRank import Rank
 from ..Model.lbsMessage import Message
-from ..IO.lbsStatistics import print_function_statistics, min_Hamming_distance
 
 
 class InformAndTransferAlgorithm(AlgorithmBase):
@@ -71,6 +70,9 @@ class InformAndTransferAlgorithm(AlgorithmBase):
         # No information about peers is known initially
         self.__known_peers = {}
 
+        # Optional parameter to conclude the iteration process iff imbalance is below the target threshold
+        self.__target_imbalance = parameters.get("target_imbalance", 0.0)
+
     def get_known_peers(self):
         """Return all known peers."""
         return self.__known_peers
@@ -207,11 +209,11 @@ class InformAndTransferAlgorithm(AlgorithmBase):
             self._logger.info(f"Iteration complete ({n_ignored} skipped ranks)")
 
             # Compute and report iteration work statistics
-            print_function_statistics(
-                self._rebalanced_phase.get_ranks(),
-                lambda x: self._work_model.compute(x),  # pylint:disable=W0108:unnecessary-lambda
-                f"iteration {i + 1} rank work",
-                self._logger)
+            stats = print_function_statistics(
+                        self._rebalanced_phase.get_ranks(),
+                        lambda x: self._work_model.compute(x),  # pylint:disable=W0108:unnecessary-lambda
+                        f"iteration {i + 1} rank work",
+                        self._logger)
 
             # Update run distributions and statistics
             self._update_distributions_and_statistics(distributions, statistics)
@@ -231,5 +233,10 @@ class InformAndTransferAlgorithm(AlgorithmBase):
                     f"Iteration {i + 1} minimum Hamming distance to optimal arrangements: {hd_min}")
                 statistics["minimum Hamming distance to optimum"].append(hd_min)
 
+            # Check if the current imbalance is within the target_imbalance range
+            if stats.statistics["imbalance"] <= self.__target_imbalance:
+                self._logger.info(f"Reached target imbalance of {self.__target_imbalance} after {i + 1} iterations.")
+                break
+
         # Report final mapping in debug mode
         self._report_final_mapping(self._logger)
diff --git a/src/lbaf/IO/lbsConfigurationValidator.py b/src/lbaf/IO/lbsConfigurationValidator.py
index 9f1ee1e..8267c64 100644
--- a/src/lbaf/IO/lbsConfigurationValidator.py
+++ b/src/lbaf/IO/lbsConfigurationValidator.py
@@ -142,6 +142,7 @@ class ConfigurationValidator:
                  "phase_id": int,
                  "parameters": {
                      "n_iterations": int,
+                     Optional("target_imbalance"): float,
                      "n_rounds": int,
                      "fanout": int,
                      "order_strategy": And(


Test patch:
diff --git a/tests/unit/config/conf_correct_clustering_target_imb.yml b/tests/unit/config/conf_correct_clustering_target_imb.yml
new file mode 100644
index 0000000..68d175a
--- /dev/null
+++ b/tests/unit/config/conf_correct_clustering_target_imb.yml
@@ -0,0 +1,41 @@
+# Specify input
+from_data:
+  data_stem: ../data/synthetic_lb_data/data
+  phase_ids:
+    - 0
+check_schema: False
+
+# Specify work model
+work_model:
+  name: AffineCombination
+  parameters:
+    alpha: 1.0
+    beta: 0.0
+    gamma: 0.0
+
+# Specify algorithm
+algorithm:
+  name: InformAndTransfer
+  phase_id: 0
+  parameters:
+    n_iterations: 4
+    target_imbalance: 0.05
+    n_rounds: 2
+    fanout: 2
+    order_strategy: arbitrary
+    transfer_strategy: Clustering
+    criterion: Tempered
+    max_objects_per_transfer: 32
+    deterministic_transfer: true
+
+# Specify output
+output_dir: ../output
+output_file_stem: output_file
+LBAF_Viz:
+  x_ranks: 2
+  y_ranks: 2
+  z_ranks: 1
+  object_jitter: 0.5
+  rank_qoi: load
+  object_qoi: load
+  force_continuous_object_qoi: True
diff --git a/tests/unit/test_configuration_validator.py b/tests/unit/test_configuration_validator.py
index 7b8c8be..0a051b5 100644
--- a/tests/unit/test_configuration_validator.py
+++ b/tests/unit/test_configuration_validator.py
@@ -232,6 +232,11 @@ class TestConfig(unittest.TestCase):
             ConfigurationValidator(config_to_validate=configuration, logger=get_logger()).main()
         self.assertEqual(err.exception.args[0], "Should be of type 'float' and magnitude > 0.0")
 
+    def test_config_validator_correct_clustering_target_imb(self):
+        with open(os.path.join(self.config_dir, "conf_correct_clustering_target_imb.yml"), "rt", encoding="utf-8") as config_file:
+            yaml_str = config_file.read()
+            configuration = yaml.safe_load(yaml_str)
+        ConfigurationValidator(config_to_validate=configuration, logger=get_logger()).main()
 
 if __name__ == "__main__":
     unittest.main()