cleaned up code with some minor changes

1 files changed, 15 insertions, 11 deletions

diff --git a/experiments/ml.pyx b/experiments/ml.pyx
index 74e5be3..c1a5616 100644
--- a/experiments/ml.pyx
+++ b/experiments/ml.pyx
@@ -51,34 +51,38 @@ def ml(dict root_victims, dict victims, dict non_victims, DTYPE_t age,
        DTYPE_t alpha, DTYPE_t delta, DTYPE_t gamma=10):
     cdef:
         int n_roots, n_victims, n_nodes, roots, i, dist, dt, t, l
-        DTYPE_t beta, all_failures, ll, beta2
+        DTYPE_t beta, ll, beta2
         list parents, failures, successes
     n_roots, n_victims = len(root_victims), len(victims)
-    n_nodes = n_victims + len(non_victims) + n_roots
+    n_nodes = n_roots + n_victims + len(non_victims)
     cdef:
         np.ndarray[DTYPE_t] probs = np.zeros(n_victims, dtype=DTYPE)
         np.ndarray[DTYPE_t] probs_fail = np.zeros(n_victims, dtype=DTYPE)
         np.ndarray[DTYPE_t] probs_nv = np.zeros(len(non_victims), dtype=DTYPE)
+
+    # loop through victims
     for i, parents in enumerate(victims.itervalues()):
         # for each victim node i, compute the probability that all its parents
         # fail to infect it, also computes the probability that its most
         # likely parent infects it
         failures = [weight_failure(dist, dt, alpha, delta, gamma)
-                    for (dist, dt) in parents]
-        all_failures = sum(failures)
+                    for (dist, dt, w1, w2, w3) in parents]
+        probs_fail[i] = sum(failures)
         successes = [weight_success(dist, dt, alpha, delta, gamma)
-                     for (dist, dt) in parents]
-        probs[i] = max(s - failures[l] for l, s in enumerate(successes))
-        probs_fail[i] = all_failures
+                     for (dist, dt, w1, w2, w3) in parents]
+        # find parent that maximizes p/\tilde{p}
+        probs[i] = max(s - failures[l] for l, s in enumerate(successes)) 
 
+    # loop through non-victims
     for i, parents in enumerate(non_victims.itervalues()):
         # for each non victim node, compute the probability that all its
         # parents fail to infect it
         failures = [weight_failure(dist, dt, alpha, delta, gamma)
-                    for (dist, dt) in parents]
+                    for (dist, dt, w1, w2, w3) in parents]
         probs_nv[i] = sum(failures)
-    probs.sort()
-    probs = probs[::-1]
+
+    # calculate log likelihood
+    probs.sort(); probs = probs[::-1] # sort probs in descending order
     cdef:
         np.ndarray[DTYPE_t] cums = probs.cumsum()
     ll =  probs_fail.sum()
@@ -87,7 +91,7 @@ def ml(dict root_victims, dict victims, dict non_victims, DTYPE_t age,
     for i in xrange(n_victims - 1, 0, -1):
         # iterate over all victim nodes to find the optimal threshold
         roots = n_roots + n_victims - 1 - i
-        beta = 1. / (1. + exp(-probs[i]))
+        beta = 1. / (1. + exp(-probs[i]))#exp(probs[i])#
         if beta > float(roots) / age:
             break
     else: