xgi-org · maximelucas · Oct 24, 2022 · Jun 13, 2022 · Jul 26, 2022 · Oct 24, 2022
diff --git a/tests/generators/test_classic.py b/tests/generators/test_classic.py
@@ -37,15 +37,24 @@ def test_flag_complex():
 
     S = xgi.flag_complex(G)
 
-    simplices = [
-        frozenset({0, 1, 2}),
+    simplices_2 = [
         frozenset({0, 1}),
         frozenset({0, 2}),
-        frozenset({1, 2}),
         frozenset({0, 3}),
+        frozenset({1, 2}),
     ]
 
-    assert S.edges.members() == simplices
+    simplices_3 = simplices_2 + [frozenset({0, 1, 2})]
+
+    assert S.edges.members() == simplices_3
+
+    S1 = xgi.flag_complex(G, ps=[1], seed=42)
+    S2 = xgi.flag_complex(G, ps=[0.5], seed=42)
+    S3 = xgi.flag_complex(G, ps=[0], seed=42)
+
+    assert S1.edges.members() == simplices_3
+    assert S2.edges.members() == simplices_2
+    assert S3.edges.members() == simplices_2
 
 
 def test_ring_lattice():
@@ -95,3 +104,4 @@ def test_sunflower():
 
     for i in range(3, 15):
         assert len(H.nodes.memberships(i)) == 1
+
diff --git a/xgi/generators/classic.py b/xgi/generators/classic.py
@@ -5,10 +5,15 @@
 
 """
 
+from collections import defaultdict
 from itertools import combinations
 from warnings import warn
 
 import networkx as nx
+import numpy as np
+
+from ..classes import SimplicialComplex
+from ..utils import py_random_state
 
 from ..exception import XGIError
 
@@ -174,16 +179,22 @@ def star_clique(n_star, n_clique, d_max):
     return H
 
 
-def flag_complex(g, max_order=2, seed=None):
+@py_random_state("seed")
+def flag_complex(G, max_order=2, ps=None, seed=None):
     """Generate a flag (or clique) complex from a
     NetworkX graph by filling all cliques up to dimension max_order.
 
     Parameters
     ----------
-    g : Networkx Graph
+    G : Networkx Graph
 
     max_order : int
         maximal dimension of simplices to add to the output simplicial complex
+    ps: list of float
+        List of probabilities (between 0 and 1) to create a
+        hyperedge from a clique, at each order d. For example,
+        ps[0] is the probability of promoting any 3-node clique (triangle) to 
+        a 3-hyperedge.
 
     Returns
     -------
@@ -196,17 +207,29 @@ def flag_complex(g, max_order=2, seed=None):
     """
     # This import needs to happen when this function is called, not when it is
     # defined.  Otherwise, a circular import error would happen.
-    from ..classes import SimplicialComplex
 
-    nodes = g.nodes()
-    edges = list(g.edges())
+    nodes = G.nodes()
+    edges = G.edges()
 
-    # compute all triangles to fill
-    max_cliques = list(nx.find_cliques(g))
+    # compute all cliques to fill
+    max_cliques = list(nx.find_cliques(G))
 
     S = SimplicialComplex()
     S.add_nodes_from(nodes)
-    S.add_simplices_from(max_cliques, max_order=max_order)
+    S.add_simplices_from(edges)
+    if not ps: # promote all cliques
+        S.add_simplices_from(max_cliques, max_order=max_order)
+        return S
+
+    # promote cliques with a given probability
+    cliques_d = defaultdict(list)
+    for x in max_cliques:
+        cliques_d[len(x)].append(x)
+
+    for i, p in enumerate(ps[:max_order - 1]):
+        d = i + 2 # simplex order
+        cliques_d_to_add = [el for el in cliques_d[d+1] if seed.random() <= p]
+        S.add_simplices_from(cliques_d_to_add, max_order=max_order)
 
     return S