#!/usr/bin/env python from collections import defaultdict import csv import sys SLOT_SIZE = 300 def fint(a): try: return int(a) except: return -1 def counter(): count = 0 while True: yield count count += 1 node_id = counter() pairs = defaultdict(set) with open(sys.argv[1]) as f: reader = csv.reader(f) next(reader) for row in reader: _, a, _, b, _, slot = map(fint, row) p = (min(a,b), max(a,b)) pairs[p].add(slot) id_map = {} traces = [] for (a, b), slots in pairs.items(): contacts = [] first = None prev = None nslots = 0 for slot in slots: if first is None: first = prev = slot nslots = 1 elif slot - prev == 1: prev = slot nslots += 1 else: contacts.append((first, first+nslots)) first = prev = slot nslots = 1 contacts.append((first, first+nslots)) if a not in id_map: id_map[a] = next(node_id) a = id_map[a] if b not in id_map: id_map[b] = next(node_id) b = id_map[b] for c in contacts: traces.append((c[0] * SLOT_SIZE, a, b, 1)) traces.append((c[1] * SLOT_SIZE, a, b, 0)) traces.sort() duration = traces[-1][0] + SLOT_SIZE nodes = next(node_id) path = sys.argv[1].split('.csv')[0] + '_reduced.csv' with open(path, 'w') as f: writer = csv.writer(f) writer.writerow(['time', 'a', 'b', 'state']) writer.writerow([duration, -1, -1, nodes]) for trace in traces: writer.writerow(trace)