From df98874a6fece7a1988d9c27f7f938b167521f2f Mon Sep 17 00:00:00 2001
From: Tobias Kempf <kempf.tobias@web.de>
Date: Fri, 20 Dec 2024 14:18:35 +0100
Subject: [PATCH] Fix rebalance init

---
 mt-kahypar/dynamic/strategies/localFM_factor.h |  3 ++-
 mt-kahypar/dynamic/strategies/rebalance.h      | 16 ++++++++++++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/mt-kahypar/dynamic/strategies/localFM_factor.h b/mt-kahypar/dynamic/strategies/localFM_factor.h
index db9dd3786..08c112de4 100644
--- a/mt-kahypar/dynamic/strategies/localFM_factor.h
+++ b/mt-kahypar/dynamic/strategies/localFM_factor.h
@@ -59,6 +59,7 @@ namespace mt_kahypar::dyn {
           //TODO: is second reset after rebalancing necessary?
           GainCachePtr::resetGainCache(_gain_cache);
 
+          //TODO does this need to be done prior to rebalancing?
           _fm->initialize(partitioned_hypergraph);
 
           Metrics best_Metrics = {mt_kahypar::metrics::quality(*partitioned_hypergraph_s, Objective::km1),
@@ -140,7 +141,7 @@ namespace mt_kahypar::dyn {
           if (skipped_changes >= step_size) {
             skipped_changes = 0;
             step_size *= 2;
-            local_fm(hypergraph, context);
+            local_fm(hypergraph, context, nodes_to_partition);
             nodes_to_partition = parallel::scalable_vector<HypernodeID>(changes_size);
           } else {
             skipped_changes++;
diff --git a/mt-kahypar/dynamic/strategies/rebalance.h b/mt-kahypar/dynamic/strategies/rebalance.h
index 41631e960..d5f02eee1 100644
--- a/mt-kahypar/dynamic/strategies/rebalance.h
+++ b/mt-kahypar/dynamic/strategies/rebalance.h
@@ -13,6 +13,7 @@ namespace mt_kahypar::dyn {
         std::optional<ds::PartitionedHypergraph<ds::StaticHypergraph>> partitioned_hypergraph_s;
         gain_cache_t _gain_cache;
         std::unique_ptr<IRebalancer> _rebalancer;
+        std::unique_ptr<IRefiner> _fm;
 
         void repartition(ds::StaticHypergraph& hypergraph_s, Context& context) {
           context.dynamic.repartition_count++;
@@ -26,10 +27,19 @@ namespace mt_kahypar::dyn {
           _gain_cache = GainCachePtr::constructGainCache(context);
           _rebalancer = RebalancerFactory::getInstance().createObject(
                   context.refinement.rebalancer, hypergraph_s.initialNumNodes(), context, _gain_cache);
+
+          context.refinement.fm.algorithm = FMAlgorithm::kway_fm;
+          context.refinement.fm.multitry_rounds = context.dynamic.multitry_localFM;
+
+          _fm = FMFactory::getInstance().createObject(
+                  context.refinement.fm.algorithm,
+                  hypergraph_s.initialNumNodes(), hypergraph_s.initialNumEdges(), context, _gain_cache, *_rebalancer);
         }
 
         //use rebalancer to rebalance partitioned_hypergraph_s
-        void rebalance(ds::StaticHypergraph& hypergraph, Context& context) {
+        void rebalance(Context& context) {
+
+          std::cout << "Rebalancing" << std::endl;
 
           GainCachePtr::resetGainCache(_gain_cache);
 
@@ -38,6 +48,8 @@ namespace mt_kahypar::dyn {
           Metrics best_Metrics = {mt_kahypar::metrics::quality(*partitioned_hypergraph_s, Objective::km1),
                                   mt_kahypar::metrics::imbalance(*partitioned_hypergraph_s, context)};
 
+          _fm->initialize(partitioned_hypergraph);
+
           _rebalancer->refineAndOutputMoves(partitioned_hypergraph, {}, moves_by_part, best_Metrics, std::numeric_limits<double>::max());
         }
 
@@ -93,7 +105,7 @@ namespace mt_kahypar::dyn {
           }
 
           if (!metrics::isBalanced(*partitioned_hypergraph_s, context)) {
-            rebalance(hypergraph, context);
+            rebalance(context);
           }
 
           ASSERT(metrics::isBalanced(*partitioned_hypergraph_s, context));