From a21509fca5e8de9ce34d86e31fddb69dcd1846f3 Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Wed, 8 Nov 2023 18:25:41 +0100
Subject: [PATCH] Fix Leiden refinement phase

---
 cpp/src/community/detail/common_methods.cuh | 36 +++++++++++++++++++++
 cpp/src/community/detail/refine_impl.cuh    | 18 +++++++----
 2 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/cpp/src/community/detail/common_methods.cuh b/cpp/src/community/detail/common_methods.cuh
index b388ba53e81..82634234398 100644
--- a/cpp/src/community/detail/common_methods.cuh
+++ b/cpp/src/community/detail/common_methods.cuh
@@ -100,6 +100,27 @@ struct reduce_op_t {
   }
 };
 
+// a workaround for cudaErrorInvalidDeviceFunction error when device lambda is used
+template <typename vertex_t, typename weight_t>
+struct count_updown_moves_op_t {
+  bool up_down{};
+  __device__ auto operator()(thrust::tuple<vertex_t, thrust::tuple<vertex_t, weight_t>> p) const
+  {
+    vertex_t old_cluster       = thrust::get<0>(p);
+    auto new_cluster_gain_pair = thrust::get<1>(p);
+    vertex_t new_cluster       = thrust::get<0>(new_cluster_gain_pair);
+    weight_t delta_modularity  = thrust::get<1>(new_cluster_gain_pair);
+
+    auto result_assignment =
+      (delta_modularity > weight_t{0})
+        ? (((new_cluster > old_cluster) != up_down) ? old_cluster : new_cluster)
+        : old_cluster;
+
+    return (delta_modularity > weight_t{0})
+             ? (((new_cluster > old_cluster) != up_down) ? false : true)
+             : false;
+  }
+};
 // a workaround for cudaErrorInvalidDeviceFunction error when device lambda is used
 template <typename vertex_t, typename weight_t>
 struct cluster_update_op_t {
@@ -394,6 +415,21 @@ rmm::device_uvector<vertex_t> update_clustering_by_delta_modularity(
     detail::reduce_op_t<vertex_t, weight_t>{},
     cugraph::get_dataframe_buffer_begin(output_buffer));
 
+  int nr_moves = thrust::count_if(
+    handle.get_thrust_policy(),
+    thrust::make_zip_iterator(thrust::make_tuple(
+      next_clusters_v.begin(), cugraph::get_dataframe_buffer_begin(output_buffer))),
+    thrust::make_zip_iterator(
+      thrust::make_tuple(next_clusters_v.end(), cugraph::get_dataframe_buffer_end(output_buffer))),
+    detail::count_updown_moves_op_t<vertex_t, weight_t>{up_down});
+
+  if (multi_gpu) {
+    nr_moves = host_scalar_allreduce(
+      handle.get_comms(), nr_moves, raft::comms::op_t::SUM, handle.get_stream());
+  }
+
+  if (nr_moves == 0) { up_down = !up_down; }
+
   thrust::transform(handle.get_thrust_policy(),
                     next_clusters_v.begin(),
                     next_clusters_v.end(),
diff --git a/cpp/src/community/detail/refine_impl.cuh b/cpp/src/community/detail/refine_impl.cuh
index 6b6470991bb..9dbc82d2c40 100644
--- a/cpp/src/community/detail/refine_impl.cuh
+++ b/cpp/src/community/detail/refine_impl.cuh
@@ -89,8 +89,9 @@ struct leiden_key_aggregated_edge_op_t {
 
     // E(Cr, S-Cr) > ||Cr||*(||S|| -||Cr||)
     bool is_dst_leiden_cluster_well_connected =
-      dst_leiden_cut_to_louvain >
-      resolution * dst_leiden_volume * (louvain_cluster_volume - dst_leiden_volume);
+      dst_leiden_cut_to_louvain > resolution * dst_leiden_volume *
+                                    (louvain_cluster_volume - dst_leiden_volume) /
+                                    total_edge_weight;
 
     // E(v, Cr-v) - ||v||* ||Cr-v||/||V(G)||
     // aggregated_weight_to_neighboring_leiden_cluster == E(v, Cr-v)?
@@ -100,9 +101,9 @@ struct leiden_key_aggregated_edge_op_t {
       if ((louvain_of_dst_leiden_cluster == src_louvain_cluster) &&
           is_dst_leiden_cluster_well_connected) {
         mod_gain = aggregated_weight_to_neighboring_leiden_cluster -
-                   resolution * src_weighted_deg * (dst_leiden_volume - src_weighted_deg) /
-                     total_edge_weight;
-
+                   resolution * src_weighted_deg * dst_leiden_volume / total_edge_weight;
+// NOTE: Disable random moves in refinement phase for now.
+#if 0
         weight_t random_number{0.0};
         if (mod_gain > 0.0) {
           auto flat_id = uint64_t{threadIdx.x + blockIdx.x * blockDim.x};
@@ -117,6 +118,8 @@ struct leiden_key_aggregated_edge_op_t {
                      ? __expf(static_cast<float>((2.0 * mod_gain) / (theta * total_edge_weight))) *
                          random_number
                      : -1.0;
+#endif
+        mod_gain = mod_gain > 0.0 ? mod_gain : -1.0;
       }
     }
 
@@ -240,11 +243,12 @@ refine_clustering(
                     wcut_deg_and_cluster_vol_triple_begin,
                     wcut_deg_and_cluster_vol_triple_end,
                     singleton_and_connected_flags.begin(),
-                    [resolution] __device__(auto wcut_wdeg_and_louvain_volume) {
+                    [resolution, total_edge_weight] __device__(auto wcut_wdeg_and_louvain_volume) {
                       auto wcut           = thrust::get<0>(wcut_wdeg_and_louvain_volume);
                       auto wdeg           = thrust::get<1>(wcut_wdeg_and_louvain_volume);
                       auto louvain_volume = thrust::get<2>(wcut_wdeg_and_louvain_volume);
-                      return wcut > (resolution * wdeg * (louvain_volume - wdeg));
+                      return wcut >
+                             (resolution * wdeg * (louvain_volume - wdeg) / total_edge_weight);
                     });
 
   edge_src_property_t<GraphViewType, weight_t> src_louvain_cluster_weight_cache(handle);