vehLanes_lxml.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys

from lxml import etree
from collections import defaultdict


def trackLanes(netstate, out):
    # veh_id -> values
    laneTimes = defaultdict(list)
    laneChanges = defaultdict(lambda: 0)
    lastEdge = defaultdict(lambda: None)
    arrivals = {}
    running = set()

    with open(out, 'w') as f:
        f.write("<vehLanes>\n")

        context = etree.iterparse(netstate, events=('end',), tag='timestep')
        for event, timestep in context:
            seen = set()
            for edges in timestep:
                for lanes in edges:
                    for vehicle in lanes:
                        seen.add(vehicle.get('id'))
                        if vehicle.get('id') not in running or laneTimes[vehicle.get('id')][-1][1] != lanes.get('id'):
                            laneTimes[vehicle.get('id')].append(
                                (timestep.get('time'), lanes.get('id')))
                            running.add(vehicle.get('id'))
                            if lastEdge[vehicle.get('id')] == edges.get('id'):
                                laneChanges[vehicle.get('id')] += 1
                            lastEdge[vehicle.get('id')] = edges.get('id')
            for veh_id in running:
                if veh_id not in seen:
                    arrivals[veh_id] = timestep.get('time')
            running = running - set(arrivals.keys())

            timestep.clear()
            while timestep.getprevious() is not None:
                del timestep.getparent()[0]

        for veh_id, lt in laneTimes.items():
            f.write('    <vehicle id="%s" laneTimes="%s" arrival="%s" laneChanges="%s"/>\n' % (
                veh_id,
                ' '.join(['%s,%s' % (t, l) for t, l in lt]),
                arrivals.get(veh_id),
                laneChanges[veh_id]))
        f.write("</vehLanes>\n")


if __name__ == "__main__":
    if len(sys.argv) != 3:
        sys.exit("call %s <netstate-dump> <output-file>" % sys.argv[0])
    trackLanes(*sys.argv[1:])