From d1172de0b18d9fb6aa837bba3265329398af44ea Mon Sep 17 00:00:00 2001
From: Tobias Kempf <kempf.tobias@web.de>
Date: Thu, 19 Dec 2024 17:09:27 +0100
Subject: [PATCH] Smarter Imbalance Checks

---
 mt-kahypar/dynamic/strategies/connectivity.h   |  3 +--
 mt-kahypar/dynamic/strategies/localFM.h        |  8 +++-----
 mt-kahypar/dynamic/strategies/localFM_factor.h | 16 ++++++++++++++--
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/mt-kahypar/dynamic/strategies/connectivity.h b/mt-kahypar/dynamic/strategies/connectivity.h
index d64c270df..fb31fadc4 100644
--- a/mt-kahypar/dynamic/strategies/connectivity.h
+++ b/mt-kahypar/dynamic/strategies/connectivity.h
@@ -90,8 +90,7 @@ namespace mt_kahypar::dyn {
 
 
           //check if imbalance is still within bounds else repartition
-          if (partitioned_hypergraph_s->partWeight(std::get<1>(*max_connectivity)) /
-              static_cast<double>(context.partition.perfect_balance_part_weights[std::get<1>(*max_connectivity)]) - 1.0 > context.partition.epsilon ) {
+          if (!metrics::isBalanced(*partitioned_hypergraph_s, context)) {
             repartition(hypergraph, context);
           }
           history.push_back(partition_result);
diff --git a/mt-kahypar/dynamic/strategies/localFM.h b/mt-kahypar/dynamic/strategies/localFM.h
index e17cb2fb3..8e2027b68 100644
--- a/mt-kahypar/dynamic/strategies/localFM.h
+++ b/mt-kahypar/dynamic/strategies/localFM.h
@@ -46,7 +46,7 @@ namespace mt_kahypar::dyn {
 
           mt_kahypar_partitioned_hypergraph_t  partitioned_hypergraph = utils::partitioned_hg_cast(*partitioned_hypergraph_s);
 
-          if (mt_kahypar::metrics::imbalance(*partitioned_hypergraph_s, context) > context.partition.epsilon) {
+          if (!metrics::isBalanced(*partitioned_hypergraph_s, context)) {
             // use rebalancer to rebalance partitioned_hypergraph_s
             parallel::scalable_vector<parallel::scalable_vector<Move>> moves_by_part;
             Metrics best_Metrics = {mt_kahypar::metrics::quality(*partitioned_hypergraph_s, Objective::km1),
@@ -127,10 +127,8 @@ namespace mt_kahypar::dyn {
 
           local_fm(hypergraph, context, local_fm_nodes);
 
-          //check if imbalance is still within bounds else repartition
-          if (mt_kahypar::metrics::imbalance(*partitioned_hypergraph_s, context) > context.partition.epsilon + 0.0000001) {
-            repartition(hypergraph, context);
-          }
+          ASSERT(metrics::isBalanced(*partitioned_hypergraph_s, context));
+
           history.push_back(partition_result);
         }
 
diff --git a/mt-kahypar/dynamic/strategies/localFM_factor.h b/mt-kahypar/dynamic/strategies/localFM_factor.h
index 87d033805..38492c96f 100644
--- a/mt-kahypar/dynamic/strategies/localFM_factor.h
+++ b/mt-kahypar/dynamic/strategies/localFM_factor.h
@@ -42,18 +42,30 @@ namespace mt_kahypar::dyn {
         }
 
         //use local_fm to refine partitioned_hypergraph_s
-        void local_fm(ds::StaticHypergraph& hypergraph, Context& context) {
+        void local_fm(ds::StaticHypergraph& hypergraph, Context& context, parallel::scalable_vector<HypernodeID> local_fm_nodes) {
 
           GainCachePtr::resetGainCache(_gain_cache);
 
           mt_kahypar_partitioned_hypergraph_t  partitioned_hypergraph = utils::partitioned_hg_cast(*partitioned_hypergraph_s);
 
+          if (!metrics::isBalanced(*partitioned_hypergraph_s, context)) {
+            // use rebalancer to rebalance partitioned_hypergraph_s
+            parallel::scalable_vector<parallel::scalable_vector<Move>> moves_by_part;
+            Metrics best_Metrics = {mt_kahypar::metrics::quality(*partitioned_hypergraph_s, Objective::km1),
+                                    mt_kahypar::metrics::imbalance(*partitioned_hypergraph_s, context)};
+            _rebalancer->refineAndOutputMoves(partitioned_hypergraph, {}, moves_by_part, best_Metrics, std::numeric_limits<double>::max());
+          }
+
+          //TODO: is second reset after rebalancing necessary?
+          GainCachePtr::resetGainCache(_gain_cache);
+
           _fm->initialize(partitioned_hypergraph);
 
           Metrics best_Metrics = {mt_kahypar::metrics::quality(*partitioned_hypergraph_s, Objective::km1),
                                   mt_kahypar::metrics::imbalance(*partitioned_hypergraph_s, context)};
 
-          _fm->refine(partitioned_hypergraph, nodes_to_partition, best_Metrics, std::numeric_limits<double>::max());
+          _fm->refine(partitioned_hypergraph, local_fm_nodes, best_Metrics, std::numeric_limits<double>::max());
+
         }
 
         PartitionID add_node_to_partitioned_hypergraph(ds::StaticHypergraph& hypergraph, Context& context, const HypernodeID& hn) {