faster version of icc cascade model

1 files changed, 32 insertions, 4 deletions

diff --git a/jpa_test/cascade_creation.py b/jpa_test/cascade_creation.py
index 88798cd..cdcfd32 100644
--- a/jpa_test/cascade_creation.py
+++ b/jpa_test/cascade_creation.py
@@ -16,23 +16,51 @@ def icc_cascade(G, p_init):
         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
+        cascade.append(active)
+    return cascade
+
+
+def icc_cascade_2(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 active.any() and susceptible.any():
+        cascade.append(active)
+        active = np.exp(np.dot(G.logmat, active)) \
+                    < np.random.rand(G.number_of_nodes())
+        active = active & susceptible
+        susceptible = susceptible - active
     return cascade
 
+
 def test():
     """
     unit test
     """
-    G = nx.erdos_renyi_graph(500, .15, directed=True)
+    G = nx.erdos_renyi_graph(1000, 1, directed=True)
+    G.logmat = np.zeros((G.number_of_nodes(), G.number_of_nodes()))
     for edge in G.edges(data=True):
-        edge[2]['weight'] = .1*np.random.rand()
-
+        edge[2]['weight'] = .3*np.random.rand()
+        G.logmat[edge[0],edge[1]] = np.log(1 - edge[2]["weight"])
+    import time
+    t0 = time.time()
     print len(icc_cascade(G, p_init=.1))
+    t1 = time.time()
+    print t1 - t0
+    print len(icc_cascade_2(G, p_init=.1))
+    print time.time() - t1
 
 if __name__ == "__main__":
     test()