Update refine accordingly

2 files changed, 13 insertions, 7 deletions

diff --git a/hawkes/data.py b/hawkes/data.py
index 5f68dd1..0f6135b 100644
--- a/hawkes/data.py
+++ b/hawkes/data.py
@@ -34,8 +34,6 @@ def load_edges(filename):
         for row in reader:
             fro, to, t, dist = map(parse, [row["from"], row["to"],
                                            row["t1"], row["dist"]])
-            if dist >= 2:
-                continue
             d = edges.get(fro, dict())
             d[to] = dist
             edges[fro] = d
@@ -70,4 +68,4 @@ if __name__ == "__main__":
     nodes = load_nodes(sys.argv[1])
     edges, events = load_edges(sys.argv[2])
     event_edges = compute_event_edges(events, edges)
-    dump((nodes, edges, events, event_edges), open("data-dist1.pickle", "wb"))
+    dump((nodes, edges, events, event_edges), open("data-all.pickle", "wb"))
diff --git a/hawkes/refine.py b/hawkes/refine.py
index 0c171bc..02bbe07 100644
--- a/hawkes/refine.py
+++ b/hawkes/refine.py
@@ -15,15 +15,23 @@ def inprod(a, b):
     return tuple(float(x * y) for x, y in zip(a, b))
 
 
+def approx(x):
+    if x > 1e-10:
+        return 1 - exp(-x)
+    else:
+        return x
+
+
 def ll(lamb, alpha, mu):
     r1 = sum(log(lamb * (1 + 0.43 * sin(0.0172 * t1 + 4.36))
                  + sum(alpha / d * mu * exp(-mu * (t1 - t2))
                        for (n2, t2, d) in s))
              for ((n1, t1), s) in event_edges.iteritems())
-    r2 = sum(sum(alpha / d * (1 - exp(-mu * (nodes[n2] - t1)))
-                 for n2, d in edges[n1].iteritems() if nodes[n2] > t1)
+    r2 = sum(sum(alpha / d * approx(mu * (nodes[n2][0] - t1))
+                 for n2, d in edges[n1].iteritems()
+                 if nodes[n2][0] > t1)
              for (n1, t1) in iter_events(events))
-    r3 = lamb * sum(nodes.itervalues())
+    r3 = lamb * sum(node[1] for node in nodes.itervalues())
     return -(r1 - r2 - r3)
 
 
@@ -58,5 +66,5 @@ def refine():
 
 
 if __name__ == "__main__":
-    nodes, edges, events, event_edges = load(open("data-dist1.pickle", "rb"))
+    nodes, edges, events, event_edges = load(open("data-all.pickle", "rb"))
     refine()