Enable using weights with head_tail

cdlib/algorithms/crisp_partition.py

@@ -2470,7 +2470,9 @@ def mod_m(g_original: object, query_node: object) -> NodeClustering:
     )
 
 
-def head_tail(g_original: object, head_tail_ratio: float = 0.4) -> NodeClustering:
+def head_tail(
+    g_original: object, head_tail_ratio: float = 0.4, weight: str = None
+) -> NodeClustering:
     """
     Identifying homogeneous communities in complex networks by applying head/tail breaks on edge betweenness given its heavy-tailed distribution.
 
@@ -2482,11 +2484,14 @@ def head_tail(g_original: object, head_tail_ratio: float = 0.4) -> NodeClusterin
     ========== ======== ========
     Undirected Directed Weighted
     ========== ======== ========
-    Yes        No       No
+    Yes        No       Yes
     ========== ======== ========
 
     :param g_original: a networkx/igraph object
     :param head_tail_ratio: head/tail division rule. Float in [0,1], dafault 0.4.
+    :param weight: edge attribute holding the edge LENGTH used by the edge-betweenness
+        computation (networkx semantics: larger = farther apart). None (default) treats
+        every edge as unit length. Default None
     :return: NodeClustering object
 
     :Example:
@@ -2505,13 +2510,13 @@ def head_tail(g_original: object, head_tail_ratio: float = 0.4) -> NodeClusterin
     """
 
     g = convert_graph_formats(g_original, nx.Graph)
-    coms = HeadTail(g)
+    coms = HeadTail(g, head_tail_ratio=head_tail_ratio, weight=weight)
 
     return NodeClustering(
         coms,
         g_original,
         "head_tail",
-        method_parameters={"head_tail_ratio": head_tail_ratio},
+        method_parameters={"head_tail_ratio": head_tail_ratio, "weight": weight},
     )
 
 

cdlib/algorithms/internal/headtail.py

@@ -2,13 +2,15 @@
 import numpy as np
 
 
-def __HeadTailCommunityDetection(G, finaledgelist, head_tail_ratio=0.6):
+def __HeadTailCommunityDetection(G, finaledgelist, head_tail_ratio=0.6, weight=None):
 
     H = nx.connected_components(G)
 
     for s in H:
         subgraph = nx.subgraph(G, s)
-        result = nx.algorithms.edge_betweenness_centrality(subgraph, normalized=False)
+        result = nx.algorithms.edge_betweenness_centrality(
+            subgraph, normalized=False, weight=weight
+        )
         edges = list(result.keys())
         values = list(result.values())
         mean = np.mean(values)
@@ -31,16 +33,24 @@ def __HeadTailCommunityDetection(G, finaledgelist, head_tail_ratio=0.6):
             else:
                 Gsub = nx.Graph()
                 for edge in edgelist:
-                    Gsub.add_edge(edge[0], edge[1])
+                    # carry edge attributes (e.g. "weight") into the recursion --
+                    # a bare add_edge(u, v) would silently drop them after level 1
+                    Gsub.add_edge(
+                        edge[0], edge[1], **G.get_edge_data(edge[0], edge[1], default={})
+                    )
                 try:
-                    __HeadTailCommunityDetection(Gsub, finaledgelist, head_tail_ratio)
+                    __HeadTailCommunityDetection(
+                        Gsub, finaledgelist, head_tail_ratio, weight
+                    )
                 except:
                     pass
     return finaledgelist
 
 
-def HeadTail(G):
-    finaledgelist = __HeadTailCommunityDetection(G, [], head_tail_ratio=0.4)
+def HeadTail(G, head_tail_ratio=0.4, weight=None):
+    finaledgelist = __HeadTailCommunityDetection(
+        G, [], head_tail_ratio=head_tail_ratio, weight=weight
+    )
     g1 = nx.Graph()
     g1.add_edges_from(finaledgelist)
     coms = [list(c) for c in nx.connected_components(g1)]

cdlib/test/test_community_discovery_models.py

@@ -950,6 +950,13 @@ def test_head_tail(self):
             self.assertEqual(type(coms.communities[0]), list)
             self.assertEqual(type(coms.communities[0][0]), int)
 
+        nx.set_edge_attributes(G, values=2, name="weight")
+        coms = algorithms.head_tail(G, 0.8, weight="weight")
+        self.assertEqual(type(coms.communities), list)
+        if len(coms.communities) > 0:
+            self.assertEqual(type(coms.communities[0]), list)
+            self.assertEqual(type(coms.communities[0][0]), int)
+
     def test_core_expansion(self):
 
         G = nx.karate_club_graph()

cdlib/test/test_partitions_comparisons.py

@@ -180,10 +180,19 @@ def test_clusim(self):
         self.assertLessEqual(score.score, 2)
         self.assertGreaterEqual(score.score, 0)
 
-        score = evaluation.rmi(louvain_communities, lp_communities)
+        # RMI = MI - log(Omega(a,b)/n) is chance-corrected, so the unnormalised
+        # form (norm_type="none", the default) is unbounded in both directions:
+        # ~-0.26 for uninformative partition pairs, >1.14 for near-identical
+        # ones. Normalising caps it at 1, modulo float error. There is no lower
+        # bound to assert -- a below-chance pair is legitimately negative, and
+        # the exactly-zero case (one side collapsed to a single community) lands
+        # a few ulps below zero.
+        score = evaluation.rmi(
+            louvain_communities, lp_communities, norm_type="normalized"
+        )
 
-        self.assertLessEqual(score.score, 1)
-        self.assertGreaterEqual(score.score, 0)
+        self.assertTrue(np.isfinite(score.score))
+        self.assertLessEqual(score.score, 1 + 1e-9)
 
         score = evaluation.geometric_accuracy(louvain_communities, lp_communities)