small bug fixes

1 files changed, 7 insertions, 5 deletions

diff --git a/jpa_test/cascade_creation.py b/jpa_test/cascade_creation.py
index 2b53963..09f0c45 100644
--- a/jpa_test/cascade_creation.py
+++ b/jpa_test/cascade_creation.py
@@ -80,8 +80,8 @@ def icc_matrixvector_for_node(cascades, node):
     M = []
     for cascade in cascades:
         t_i = cascade.infection_time(node)[0]
-        if t_i > 0 or t_i is None:
-            indicator = np.zeros(len(cascade))
+        if t_i != 0:
+            indicator = np.zeros(len(cascade[:t_i]))
             if t_i > 0:
                 indicator[-1] = 1
             w.append(indicator)
@@ -107,6 +107,9 @@ def icc_cascade(G, p_init):
                     < np.random.rand(G.number_of_nodes())
         active = active & susceptible
         susceptible = susceptible - active
+    if not cascade:
+        print "Empty cascade, consider changing p_init or n_nodes. Retrying."
+        return icc_cascade(G, p_init)
     return cascade
 
 
@@ -122,12 +125,11 @@ def test():
     unit test
     """
     G = InfluenceGraph(max_proba = .3)
-    G.erdos_init(n = 100, p = 1)
+    G.erdos_init(n = 10, p = 1)
     import time
     t0 = time.time()
-    A = generate_cascades(G, .1, 10)
+    A = generate_cascades(G, .1, 2)
     M, w = icc_matrixvector_for_node(A, 0)
-    print w
     t1 = time.time()
     print t1 - t0