from csv import DictReader import sys from cPickle import dump from os.path import splitext def build_network(filename): victims = {} non_victims = {} age = 0. with open(filename) as fh: reader = DictReader(fh) for row in reader: print age from_, to = int(float(row["from"])), int(float(row["to"])) dist = int(row["dist"]) if int(float(row["dist"])) > 2: continue # 'to' is a victim if row["t2"] != "NA": dt = int(row["t2"]) - int(row["t1"]) parent = (dist, dt) if to not in victims: age += int(row["t2"]) - int(row["spawn2"]) victims[to] = [] victims[to].append(parent) if from_ not in victims: age += int(row["t1"]) - int(row["spawn1"]) victims[from_] = [] # 'to' is not a victim else: dt = 3012 - int(row["t1"]) parent = (dist, dt) if to not in non_victims: age += 3012 - int(row["spawn2"]) non_victims[to] = [] non_victims[to].append(parent) if from_ not in victims: age += int(row["t1"]) - int(row["spawn1"]) victims[from_] = [] root_victims = {} for victim in victims.keys(): if not victims[victim]: del victims[victim] root_victims[victim] = [] print age return root_victims, victims, non_victims, age if __name__ == "__main__": if len(sys.argv) < 2: sys.exit("usage: {0} ".format(sys.argv[0])) filename = sys.argv[1] root, _ = splitext(filename) root_victims, victims, non_victims, age = build_network(filename) dump((root_victims, victims, non_victims, age), open(root + ".pickle", "w"))