1 files changed, 33 insertions, 14 deletions
diff --git a/hawkes/refine.py b/hawkes/refine.py
index 675bee3..3dfcfcc 100644
--- a/hawkes/refine.py
+++ b/hawkes/refine.py
@@ -1,6 +1,8 @@
 from cPickle import load
 from itertools import product
-from math import exp, log
+from math import exp, log, sin
+import sys
+from multiprocessing import Pool
 
 
 def iter_events(events):
@@ -13,31 +15,48 @@ def inprod(a, b):
     return tuple(float(x * y) for x, y in zip(a, b))
 
 
-def ll(a):
-    lamb, alpha, mu = a
-
-    r1 = sum(log(lamb + sum(alpha * w * mu * exp(-mu * (t1 - t2))
-                            for (n2, t2, w) in s))
+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 * w * (1 - exp(-mu * (nodes[n2] - t1)))
-                 for n2, w in edges[n1].iteritems() if nodes[n2] > t1)
+    r2 = sum(sum(alpha / d * (1 - exp(-mu * (nodes[n2] - t1)))
+                 for n2, d in edges[n1].iteritems() if nodes[n2] > t1)
              for (n1, t1) in iter_events(events))
     r3 = lamb * sum(nodes.itervalues())
     return -(r1 - r2 - r3)
 
 
-if __name__ == "__main__":
-    nodes, edges, events, event_edges = load(open("data.pickle", "rb"))
+def get_values():
     d = {}
-    for line in open("test.txt"):
+    for line in open(sys.argv[1]):
         v = map(float, line.strip().split())
         d[tuple(v[:3])] = v[3]
     l = d.items()
     l.sort(key=lambda x: x[1])
-    for a, _ in l[:10]:
-        t = [1. / i for i in range(2, 6)] + [float(i) for i in range(1, 6)]
+    for line in open("refine.txt"):
+        v = map(float, line.strip().split())
+        d[tuple(v[:3])] = v[3]
+    for a, _ in l[:100]:
+        t = [1. / i for i in range(2, 4)] + [float(i) for i in range(1, 4)]
         for b in product(t, repeat=3):
             l, al, m = inprod(a, b)
             if (l, al, m) in d:
                 continue
-            print l, al, m, ll((l, al, m))
+            yield (l, al, m)
+
+
+def refine():
+    p = Pool(5)
+
+    def aux(x):
+        l, a, m = x
+        print l, a, m, ll(l, a, m)
+        sys.stdout.flush()
+
+    p.map(aux, get_values())
+
+
+if __name__ == "__main__":
+    nodes, edges, events, event_edges = load(open("data2.pickle", "rb"))
+    refine()