Add simulation code, simulating hawkes processes is fun!HEADmaster

1 files changed, 115 insertions, 0 deletions

diff --git a/hawkes_experiments/simulation.py b/hawkes_experiments/simulation.py
new file mode 100644
index 0000000..b3c1b86
--- /dev/null
+++ b/hawkes_experiments/simulation.py
@@ -0,0 +1,115 @@
+from numpy.random import poisson, uniform
+import numpy as np
+from math import cos, exp, log, pi
+from itertools import chain
+from networkx.generators.random_graphs import watts_strogatz_graph
+from csv import DictWriter
+import networkx as nx
+
+
+def homogenous_poisson_times(tmax, mu):
+    n = poisson(mu * tmax)
+    u = np.sort(uniform(size=n))
+    return tmax * u
+
+
+def _ih_poisson_times(tmax, integral, inverse):
+    itmax = integral(tmax)
+    n = poisson(itmax)
+    if n == 0:
+        return np.array([])
+    u = uniform(size=n)
+    iv = np.vectorize(lambda x: inverse(x * itmax))
+    u = np.sort(iv(u))
+    return u
+
+
+def sin_rate_times(tmax, mu, a, w, phi):
+    def integral(t):
+        return mu * t + mu * a / w * (cos(phi) - cos(w * t + phi))
+
+    def inverse(z):
+        a, b = 0, tmax
+        while b - a >= 1e-10:
+            mid = (a + b) / 2.
+            if integral(mid) <= z:
+                a = mid
+            else:
+                b = mid
+        return mid
+
+    return _ih_poisson_times(tmax, integral, inverse)
+
+
+def exp_rate_times(tmax, alpha, beta):
+    def integral(t):
+        return alpha * (1 - exp(-beta * t))
+
+    def inverse(z):
+        return - 1 / beta * log(1 - z / alpha)
+
+    return _ih_poisson_times(tmax, integral, inverse)
+
+
+def bfs(g, node):
+    d = {node: 0}
+    layer = g[node]
+    dist = 0
+    while dist < 3:
+        dist += 1
+        for n in layer:
+            d[n] = dist
+        layer = chain.from_iterable([u for u in g[n] if u not in d]
+                                    for n in layer)
+    del d[node]
+    return d
+
+
+def simulate(g, tmax, mu, a, w, phi, alpha, beta):
+    events = set()
+    gen = set()
+    gen_number = 1
+
+    for node in xrange(len(g)):
+        gen |= set((node, t) for t in sin_rate_times(tmax, mu, a, w, phi))
+        events |= gen
+
+    print gen_number, len(events)
+
+    while gen:
+        gen_number += 1
+        next_gen = set()
+        for node, time in gen:
+            for neighbor, d in bfs(g, node).items():
+                next_gen |= set((neighbor, time + t)
+                                for t in exp_rate_times(tmax - time,
+                                                        alpha / d ** 2, beta))
+        gen = next_gen
+        events |= gen
+        print gen_number, len(events)
+    return events
+
+
+def write(g, events):
+    with open('test_nodes.csv', 'w') as fh:
+        fieldnames = ['name', 'fatal_day']
+        writer = DictWriter(fh, fieldnames=fieldnames)
+        writer.writeheader()
+        for n in g:
+            writer.writerow({'name': n, 'fatal_day': 'NA'})
+    with open('test_edges.csv', 'w') as fh:
+        fieldnames = ['from', 'to', 't1', 'dist']
+        writer = DictWriter(fh, fieldnames=fieldnames)
+        writer.writeheader()
+        for n, t in events:
+            for neighbor, d in bfs(g, n).items():
+                writer.writerow({'from': n, 'to': neighbor,
+                                 'dist': d, 't1': t})
+
+
+if __name__ == "__main__":
+    g = watts_strogatz_graph(10000, 30, 0.2)
+    print nx.is_connected(g)
+    events = simulate(g, 3012, 1.19e-5, 0.43, 2 * pi / 365.24, 4.36,
+                      7.82e-3, 3.74e-3)
+    write(g, events)