Correct likelihood computation, still have to write the optimization code

1 files changed, 35 insertions, 22 deletions

diff --git a/simulation/main.py b/simulation/main.py
index 178700a..63c8c02 100644
--- a/simulation/main.py
+++ b/simulation/main.py
@@ -2,38 +2,51 @@ import numpy as np
 import numpy.random as nr
 
 
-def likelihood(p, cascade, node):
-    m = build_matrix(cascade, node)
-    x = m[:-1, :]
-    y = m[1:, node].reshape(-1)
+def likelihood(p, cascades, node):
+    x, y = build_matrix(cascades, node)
     a = 1 - np.exp(-np.dot(x, p))
-    return y * np.log(a) + (1 - y) * np.log(1-a)
+    return np.inner(y, np.log(a)) + np.inner((1 - y), np.log(1-a))
 
 
-def build_matrix(cascade, node):
-    try:
-        m = np.vstack(x for x, s in cascade if s[node])
-    except ValueError:
-        return None
-    return m
+def build_matrix(cascades, node):
 
+    def aux(cascade, node):
+        try:
+            m = np.vstack(x for x, s in cascade if s[node])
+        except ValueError:
+            return None
+        x = m[:-1, :]
+        y = m[1:, node]
+        return x, y
 
-def simulate_cascade(g, x):
-    s = x ^ np.ones(g.shape[0]).astype(bool)
-    yield x, s
+    pairs = (aux(cascade, node) for cascade in cascades)
+    xs, ys = zip(*(pair for pair in pairs if pair))
+    x = np.vstack(xs)
+    y = np.concatenate(ys)
+    return x, y
+
+
+def simulate_cascade(x, graph):
+    susc = x ^ np.ones(graph.shape[0]).astype(bool)
+    yield x, susc
     while np.any(x):
-        x = 1 - np.exp(-np.dot(g.T, x))
+        x = 1 - np.exp(-np.dot(graph.T, x))
         y = nr.random(x.shape[0])
-        x = (x >= y) & s
-        yield x, s
-        s = x ^ s
+        x = (x >= y) & susc
+        yield x, susc
+        susc ^= x
+
+
+def simulate_cascades(n, graph):
+    for _ in xrange(n):
+        x = np.zeros(g.shape[0]).astype(bool)
+        x[nr.random()] = True
+        yield simulate_cascade(x, graph)
 
 
 if __name__ == "__main__":
     g = np.array([[0, 1, 1, 0], [1, 0, 0, 1], [1, 0, 0, 1], [0, 1, 1, 0]])
     p = 0.5
     g = np.log(1. / (1 - p * g))
-    x = np.zeros(g.shape[0]).astype(bool)
-    x[0] = 1
-    cascade = simulate_cascade(g, x)
-    print likelihood(np.ones(g.shape[0]) * 0.5, cascade, 1)
+    cascades = simulate_cascades(10, g)
+    print likelihood(p * np.ones(g.shape[0]), cascades, 3)