Skip to content
Snippets Groups Projects
shed_reduce 1.56 KiB
Newer Older
  • Learn to ignore specific revisions
  • Jarrod Pas's avatar
    Jarrod Pas committed
    #!/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)