1 files changed, 109 insertions, 0 deletions
diff --git a/ic_experiments/ml.pyx b/ic_experiments/ml.pyx
new file mode 100644
index 0000000..e1bf26b
--- /dev/null
+++ b/ic_experiments/ml.pyx
@@ -0,0 +1,109 @@
+# cython: boundscheck=False, cdivision=True
+import numpy as np
+cimport numpy as np
+from libc.math cimport log, exp
+
+DTYPE = np.float64
+ctypedef np.float_t DTYPE_t
+
+cdef DTYPE_t weight_success(int dist, int dt, DTYPE_t alpha, DTYPE_t delta, DTYPE_t lmbda,
+                            DTYPE_t t_scale, DTYPE_t w1, DTYPE_t w2, DTYPE_t w3):
+    """weight for successful infection, exponential time model"""
+    cdef DTYPE_t structural, temporal, result
+    # structural = delta ** (dist**1)
+    structural = delta * lmbda**(dist-1)
+    temporal = log(exp(alpha/t_scale)-1.) - alpha*dt/t_scale
+    # temporal = 1. / (1. + (dt - 1.)/alpha)**0.01 - 1. / (1. + dt/alpha)**0.01
+    result = log(structural) + temporal
+    return result
+
+cdef DTYPE_t weight_failure(int dist, int dt, DTYPE_t alpha, DTYPE_t delta, DTYPE_t lmbda,
+                            DTYPE_t t_scale, DTYPE_t w1, DTYPE_t w2, DTYPE_t w3):
+    """weight for failed infection, exponential time model"""
+    cdef DTYPE_t structural, temporal, result
+    # structural = delta ** (dist**1)
+    structural = delta * lmbda**(dist-1)
+    temporal = exp(-alpha * dt/t_scale)
+    # temporal = 1.  - 1. / (1. + dt/alpha)**0.01
+    result = log(1. - structural + structural * temporal)
+    return result
+
+def ml(dict root_victims, dict victims, dict non_victims,
+       DTYPE_t alpha, DTYPE_t delta, DTYPE_t lmbda):
+    cdef:
+        int n_roots, n_victims, roots, i, dist, dt, t, l, n_days
+        DTYPE_t beta, ll, beta_add, max_beta, max_beta_add
+        list parents, failures, successes
+    n_roots, n_victims = len(root_victims), len(victims)
+    n_nodes = 148152
+    n_days = 3012
+    t_scale = 1.
+    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] parent_dists = np.zeros(n_victims, dtype=DTYPE)
+        np.ndarray[DTYPE_t] parent_dts = 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, lmbda, t_scale, w1, w2, w3)
+                    for (prnt, dist, dt, w1, w2, w3) in parents]
+        probs_fail[i] = sum(failures)
+        successes = [weight_success(dist, dt, alpha, delta, lmbda, t_scale, w1, w2, w3)
+                     for (prnt, dist, dt, w1, w2, w3) in parents]
+        dists = [dist for (prnt, dist, dt, w1, w2, w3) in parents]
+        dts = [dt for (prnt, dist, dt, w1, w2, w3) in parents]
+        # find parent that maximizes log(p) - log(\tilde{p})
+        probs[i] = float("-inf")
+        for l, s in enumerate(successes):
+            prob = s - failures[l]
+            if prob > probs[i]:
+                probs[i] = prob
+                parent_dists[i] = dists[l]
+                parent_dts[i] = dts[l]
+                # probs_fail[i] = failures[l]
+
+    # 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, lmbda, t_scale, w1, w2, w3)
+                    for (prnt, dist, dt, w1, w2, w3) in parents]
+        probs_nv[i] = sum(failures)
+
+    # calculate log likelihood
+    ll =  probs_fail.sum() # add probability that all edges to all victims fail
+    ll += probs_nv.sum() # add probability that all edges to non_victims fail
+
+    max_beta_add = float('-inf')
+    # iterate over all victim nodes to find the optimal threshold
+    for beta in np.arange(0.01, 1., 0.01):
+        thresh = log(beta/(n_days*(1.-beta)))
+        seeds = probs<thresh
+        non_seeds = probs>=thresh
+        roots = n_roots + sum(seeds)
+
+        beta_add = 0.
+        # add probability for realized edges and subtract probability these edges fail
+        beta_add += (probs[non_seeds]).sum()
+        # add probability for the seeds and non-seeds
+        beta_add += roots * log(beta/n_days) + (n_nodes-roots) * log(1. - beta)
+
+        if beta_add > max_beta_add:
+            max_beta = beta
+            max_roots = roots
+            max_beta_add = beta_add
+            pdists = (parent_dists[non_seeds]).mean()
+            pdts = (parent_dts[non_seeds]).mean()
+            # print 'beta:', max_beta, 'add:', max_beta_add, 'roots:', max_roots
+
+    ll += max_beta_add
+    roots = max_roots
+    beta = max_beta
+    print 'dist:', pdists
+    print 'dt:', pdts
+    return (beta, roots, ll)