1 files changed, 39 insertions, 0 deletions
diff --git a/simulation/main.py b/simulation/main.py
new file mode 100644
index 0000000..178700a
--- /dev/null
+++ b/simulation/main.py
@@ -0,0 +1,39 @@
+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)
+    a = 1 - np.exp(-np.dot(x, p))
+    return y * np.log(a) + (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 simulate_cascade(g, x):
+    s = x ^ np.ones(g.shape[0]).astype(bool)
+    yield x, s
+    while np.any(x):
+        x = 1 - np.exp(-np.dot(g.T, x))
+        y = nr.random(x.shape[0])
+        x = (x >= y) & s
+        yield x, s
+        s = x ^ s
+
+
+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)