1 files changed, 19 insertions, 4 deletions
diff --git a/simulation/main.py b/simulation/main.py
index 1c0b3e8..a916034 100644
--- a/simulation/main.py
+++ b/simulation/main.py
@@ -12,6 +12,13 @@ from six.moves import range
 seaborn.set_style("white")
 
 
+def create_random_graph(n_nodes, p=.5):
+    graph = .5 * np.random.binomial(2, p=.5, size=(n_nodes, n_nodes))
+    for k in range(len(graph)):
+        graph[k, k] = 0
+    return np.log(1. / (1 - p * graph))
+
+
 def simulate_cascade(x, graph):
     """
     Simulate an IC cascade given a graph and initial state.
@@ -22,7 +29,6 @@ def simulate_cascade(x, graph):
     """
     yield x, np.zeros(graph.shape[0], dtype=bool)
     susc = np.ones(graph.shape[0], dtype=bool)
-    #yield x, susc
     while np.any(x):
         susc = susc ^ x  # nodes infected at previous step are now inactive
         if not np.any(susc):
@@ -39,6 +45,14 @@ def uniform_source(graph, *args, **kwargs):
     return x0
 
 
+def var_source(graph, t, node_p=None, *args, **kwargs):
+    if node_p is None:
+        node_p = np.ones(len(graph)) / len(graph)
+    x0 = np.zeros(graph.shape[0], dtype=bool)
+    x0[nr.choice(a=len(graph), p=node_p)] = True
+    return x0
+
+
 def simulate_cascades(n, graph, source=uniform_source):
     for t in range(n):
         x0 = source(graph, t)
@@ -58,9 +72,10 @@ def build_cascade_list(cascades, collapse=False):
 
 
 if __name__ == "__main__":
-    g = np.array([[0, 0, 1], [0, 0, 0.5], [0, 0, 0]])
-    p = 0.5
-    g = np.log(1. / (1 - p * g))
+    #g = np.array([[0, 0, 1], [0, 0, 0.5], [0, 0, 0]])
+    #p = 0.5
+    #g = np.log(1. / (1 - p * g))
+    g = create_random_graph(n_nodes=3)
     print(g)
     sizes = [10**3]
     for si in sizes: