Add optimization code

1 files changed, 16 insertions, 2 deletions

diff --git a/hawkes/sanity.py b/hawkes/sanity.py
index c837757..529203c 100644
--- a/hawkes/sanity.py
+++ b/hawkes/sanity.py
@@ -1,5 +1,19 @@
 from cPickle import load
+from csv import reader
+from data2 import parse
+import sys
+
+
+def parse_row(row):
+    return set(e for e in map(parse, row[2:]) if e)
+
 
 if __name__ == "__main__":
-    nodes, edges, events, event_edges = load(open("data.pickle", "rb"))
-    print len(nodes)
+    nodes, edges, events, event_edges = load(open("data2.pickle", "rb"))
+    with open(sys.argv[1]) as fh:
+        fh.readline()
+        reader = reader(fh)
+        d = {parse(row[1]): parse_row(row) for row in reader}
+    d = {k: v for (k, v) in d.iteritems() if v}
+    for k in d:
+        print len(d[k])