Enable weights handling and exposes deg_corr and rec_type parameters …

cdlib/algorithms/crisp_partition.py

@@ -1586,8 +1586,24 @@ def principled_clustering(
     )
 
 
+def __prepare_sbm_weights(g_nx: object, weights) -> object:
+    """Resolve the SBM ``weights`` argument (edge-attribute name or list) to an edge-attribute name."""
+    if weights is None or isinstance(weights, str):
+        return weights
+    weights = list(weights)
+    n_edges = g_nx.number_of_edges()
+    if len(weights) != n_edges:
+        raise ValueError(f"Got {len(weights)} weights but {n_edges} edges.")
+    attr = "__cdlib_sbm_weight__"
+    nx.set_edge_attributes(g_nx, dict(zip(g_nx.edges(), weights)), attr)
+    return attr
+
+
 def sbm_dl(
     g_original: object,
+    weights: object = None,
+    rec_type: str = "real-normal",
+    deg_corr: bool = True,
 ) -> NodeClustering:
     """Efficient Monte Carlo and greedy heuristic for the inference of stochastic block models.
 
@@ -1599,10 +1615,13 @@ def sbm_dl(
     ========== ======== ========
     Undirected Directed Weighted
     ========== ======== ========
-    Yes        No       No
+    Yes        No       Yes
     ========== ======== ========
 
     :param g_original: network/igraph object
+    :param weights: list of double, or edge attribute. Fits a weighted SBM when given. Default None.
+    :param rec_type: graph-tool edge covariate type (e.g. "real-normal", "discrete-poisson"). Default "real-normal".
+    :param deg_corr: fit a degree-corrected SBM. Default True.
     :return: NodeClustering object
 
 
@@ -1631,18 +1650,35 @@ def sbm_dl(
             "(apt-get install python3-graph-tool, brew install graph-tool, etc.)"
         )
     gt_g = convert_graph_formats(g_original, nx.Graph)
-    gt_g, label_map = __from_nx_to_graph_tool(gt_g)
-    state = gt.minimize_blockmodel_dl(gt_g)
+    weight_attr = __prepare_sbm_weights(gt_g, weights)
+    gt_g, label_map = __from_nx_to_graph_tool(gt_g, weight=weight_attr)
+    state_args = {"deg_corr": deg_corr}
+    if weight_attr is not None:
+        state_args["recs"] = [gt_g.ep[weight_attr]]
+        state_args["rec_types"] = [rec_type]
+    state = gt.minimize_blockmodel_dl(gt_g, state_args=state_args)
 
     affiliations = state.get_blocks().get_array()
     affiliations = {label_map[i]: affiliations[i] for i in range(len(affiliations))}
     coms = affiliations2nodesets(affiliations)
     coms = [list(v) for k, v in coms.items()]
-    return NodeClustering(coms, g_original, "SBM", method_parameters={})
+    return NodeClustering(
+        coms,
+        g_original,
+        "SBM",
+        method_parameters={
+            "weights": weights,
+            "rec_type": rec_type if weights is not None else None,
+            "deg_corr": deg_corr,
+        },
+    )
 
 
 def sbm_dl_nested(
     g_original: object,
+    weights: object = None,
+    rec_type: str = "real-normal",
+    deg_corr: bool = True,
 ) -> NodeClustering:
     """Efficient Monte Carlo and greedy heuristic for the inference of stochastic block models. (nested)
 
@@ -1655,10 +1691,13 @@ def sbm_dl_nested(
     ========== ======== ========
     Undirected Directed Weighted
     ========== ======== ========
-    Yes        No       No
+    Yes        No       Yes
     ========== ======== ========
 
     :param g_original: igraph/networkx object
+    :param weights: list of double, or edge attribute. Fits a weighted SBM when given. Default None.
+    :param rec_type: graph-tool edge covariate type (e.g. "real-normal", "discrete-poisson"). Default "real-normal".
+    :param deg_corr: fit a degree-corrected SBM. Default True.
     :return: NodeClustering object
 
 
@@ -1692,16 +1731,30 @@ def sbm_dl_nested(
             )
 
     gt_g = convert_graph_formats(g_original, nx.Graph)
-    gt_g, label_map = __from_nx_to_graph_tool(gt_g)
-    state = gt.minimize_nested_blockmodel_dl(gt_g)
+    weight_attr = __prepare_sbm_weights(gt_g, weights)
+    gt_g, label_map = __from_nx_to_graph_tool(gt_g, weight=weight_attr)
+    state_args = {"deg_corr": deg_corr}
+    if weight_attr is not None:
+        state_args["recs"] = [gt_g.ep[weight_attr]]
+        state_args["rec_types"] = [rec_type]
+    state = gt.minimize_nested_blockmodel_dl(gt_g, state_args=state_args)
 
     level0 = state.get_levels()[0]
 
     affiliations = level0.get_blocks().get_array()
     affiliations = {label_map[i]: affiliations[i] for i in range(len(affiliations))}
     coms = affiliations2nodesets(affiliations)
     coms = [list(v) for k, v in coms.items()]
-    return NodeClustering(coms, g_original, "SBM_nested", method_parameters={})
+    return NodeClustering(
+        coms,
+        g_original,
+        "SBM_nested",
+        method_parameters={
+            "weights": weights,
+            "rec_type": rec_type if weights is not None else None,
+            "deg_corr": deg_corr,
+        },
+    )
 
 
 def markov_clustering(

cdlib/test/test_community_discovery_models.py

@@ -516,6 +516,28 @@ def test_sbm_nested_dl(self):
                 self.assertEqual(type(coms.communities[0]), list)
                 self.assertEqual(type(coms.communities[0][0]), str)
 
+    def test_sbm_dl_weighted(self):
+        if gt is not None:
+            g = get_string_graph()
+            for u, v in g.edges():
+                g[u][v]["weight"] = 1.0
+            weight_list = [g[u][v]["weight"] for u, v in g.edges()]
+            for coms in (
+                algorithms.sbm_dl(g, weights="weight"),
+                algorithms.sbm_dl(g, weights=weight_list),
+                algorithms.sbm_dl(g, weights="weight", rec_type="real-exponential"),
+                algorithms.sbm_dl(g, weights="weight", deg_corr=False),
+                algorithms.sbm_dl_nested(g, weights="weight"),
+                algorithms.sbm_dl_nested(g, weights=weight_list),
+            ):
+                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]), str)
+            # a mismatched weight list must be rejected
+            with self.assertRaises(ValueError):
+                algorithms.sbm_dl(g, weights=[1.0])
+
     # def test_danmf(self):
     #    if karateclub is None:
     #        return

cdlib/utils.py

@@ -33,11 +33,12 @@ def suppress_stdout():
             sys.stderr = old_stderr
 
 
-def __from_nx_to_graph_tool(g: object, directed: bool = None) -> object:
+def __from_nx_to_graph_tool(g: object, directed: bool = None, weight: str = None) -> object:
     """
 
     :param g:
     :param directed:
+    :param weight: edge attribute to carry as a graph-tool edge property. Default None.
     :return:
     """
 
@@ -54,7 +55,15 @@ def __from_nx_to_graph_tool(g: object, directed: bool = None) -> object:
     node_map = {v: i for i, v in enumerate(g.nodes())}
 
     gt_g.add_vertex(len(node_map))
-    gt_g.add_edge_list([(node_map[u], node_map[v]) for u, v in g.edges()])
+
+    if weight is None:
+        gt_g.add_edge_list([(node_map[u], node_map[v]) for u, v in g.edges()])
+    else:
+        # fill the weight property positionally as edges are added
+        eprop = gt_g.new_edge_property("double")
+        edges = [(node_map[u], node_map[v], d.get(weight, 1.0)) for u, v, d in g.edges(data=True)]
+        gt_g.add_edge_list(edges, eprops=[eprop])
+        gt_g.edge_properties[weight] = eprop
 
     return gt_g, {v: k for k, v in node_map.items()}