messy cascade_creation.py

1 files changed, 38 insertions, 0 deletions

diff --git a/jpa_test/cascade_creation.py b/jpa_test/cascade_creation.py
new file mode 100644
index 0000000..88798cd
--- /dev/null
+++ b/jpa_test/cascade_creation.py
@@ -0,0 +1,38 @@
+import networkx as nx
+import numpy as np
+
+def icc_cascade(G, p_init):
+    """
+    input: graph with prob as edge attr
+    returns: 2D boolean matrix with indep. casc.
+        where True means node was active at that time step
+    p_init: proba that node in seed set
+    """
+    susceptible = np.ones(G.number_of_nodes(), dtype=bool)
+    active = np.random.rand(G.number_of_nodes()) < p_init
+    susceptible = susceptible - active
+    cascade = []
+    while sum(active) and sum(susceptible):
+        cascade.append(active)
+        tmp = np.zeros(G.number_of_nodes(), dtype=bool)
+        for node in np.where(active)[0]:
+            #rand_p = np.random.rand(G.degree(2))
+            for edge in G.edges(node, data=True):
+                tmp[edge[1]] += np.random.rand() < edge[2]["weight"] \
+                                and susceptible[edge[1]]
+        active = tmp
+        susceptible = susceptible - active
+    return cascade
+
+def test():
+    """
+    unit test
+    """
+    G = nx.erdos_renyi_graph(500, .15, directed=True)
+    for edge in G.edges(data=True):
+        edge[2]['weight'] = .1*np.random.rand()
+
+    print len(icc_cascade(G, p_init=.1))
+
+if __name__ == "__main__":
+    test()