1 files changed, 38 insertions, 25 deletions
diff --git a/jpa_test/cascade_creation.py b/jpa_test/cascade_creation.py
index cc84c70..dbb9e54 100644
--- a/jpa_test/cascade_creation.py
+++ b/jpa_test/cascade_creation.py
@@ -70,6 +70,35 @@ class Cascade(list):
         return candidate_infectors
 
 
+def icc_cascade(G, p_init):
+    """
+    Returns boolean vectors for one cascade
+    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 & np.logical_not(active)
+    cascade = Cascade()
+    while active.any():
+        cascade.append(active)
+        active = np.exp(np.dot(G.logmat, active)) \
+                    < np.random.rand(G.number_of_nodes())
+        active = active & susceptible
+        susceptible = susceptible & np.logical_not(active)
+    if not cascade:
+        print "Empty cascade, consider changing p_init or n_nodes. Retrying."
+        return icc_cascade(G, p_init)
+    return cascade
+
+
+def generate_cascades(G, p_init, n_cascades):
+    """
+    returns list of cascades
+    """""
+    return [icc_cascade(G,p_init) for i in xrange(n_cascades)]
+
+
 def icc_matrixvector_for_node(cascades, node):
     """
     for the ICC model:
@@ -103,33 +132,17 @@ def normalize_matrix(M):
     return normalize(M.astype("float32"), axis=0, norm="l2")
 
 
-def icc_cascade(G, p_init):
+def add_edges_from_proba_vector(G, p_node, node, floor_cstt):
     """
-    Returns boolean vectors for one cascade
-    True means node was active at that time step
-    p_init: proba that node in seed set
+    Takes proba vector, node and adds edges to G by flooring very small
+    probabilities
+    Also updates G's mat matrix
     """
-    susceptible = np.ones(G.number_of_nodes(), dtype=bool)
-    active = np.random.rand(G.number_of_nodes()) < p_init
-    susceptible = susceptible & np.logical_not(active)
-    cascade = Cascade()
-    while active.any():
-        cascade.append(active)
-        active = np.exp(np.dot(G.logmat, active)) \
-                    < np.random.rand(G.number_of_nodes())
-        active = active & susceptible
-        susceptible = susceptible & np.logical_not(active)
-    if not cascade:
-        print "Empty cascade, consider changing p_init or n_nodes. Retrying."
-        return icc_cascade(G, p_init)
-    return cascade
-
-
-def generate_cascades(G, p_init, n_cascades):
-    """
-    returns list of cascades
-    """""
-    return [icc_cascade(G,p_init) for i in xrange(n_cascades)]
+    floor_parent = np.nonzero(p_node*(p_node > floor_cstt))
+    for parent in floor_parent[0]:
+        G.add_edge(parent, node)
+    #TODO: update G's mat matrix
+    return G
 
 
 def test():