matsim_signals_gen/signals_gen.py

#!/usr/bin/env python3
''' generates signals xml files for MATSim from OSM data '''

import lxml.etree as etree

INPUT_NETWORK = 'input/network.xml'
INPUT_SIGNALS = 'input/signals.xml'
OUT_SIGNAL_SYSTEMS = 'out_signals_systems.xml'
OUT_SIGNAL_GROUPS = 'out_signals_groups.xml'
OUT_SIGNAL_CONTROL = 'out_signals_control.xml'
CYCLE_TIME = 120

def make_subelement(parent_node, sub_name, sub_attrs=None):
    ''' creates an xml sub element and set its attributes '''
    sub = etree.SubElement(parent_node, sub_name)
    if sub_attrs is None:
        return sub
    for attr, value in sub_attrs.items():
        sub.set(attr, value)
    return sub

def make_signal_systems(network_root, signals):
    ''' make signal systems xml root '''
    signal_systems = etree.Element('signalSystems')
    for n_sig, sig in enumerate(signals):
        sig_node_id = sig.get('id')
        links = network_root.xpath("/network/links/link[@to='{}']".format(sig_node_id))
        if links:
            signal_system = make_subelement(signal_systems, 'signalSystem', {'id': str(n_sig+1)})
            signals = make_subelement(signal_system, 'signals')
            for n_link, link in enumerate(links):
                make_subelement(signals, 'signal', {'linkIdRef': link.get('id'), 'id': str(n_link+1)})
    return signal_systems

def make_signal_groups(signal_systems):
    ''' make signal groups xml root '''
    signal_groups = etree.Element('signalGroups')
    for n_sig_sys, sig_sys in enumerate(signal_systems.xpath('/signalSystems/signalSystem')):
        signal_system = make_subelement(signal_groups, 'signalSystem', {'refId': sig_sys.get('id')})
        for n_sig, sig in enumerate(sig_sys.xpath('signals/signal')):
            signal_group = make_subelement(signal_system, 'signalGroup', {'id': str(n_sig+1)})
            make_subelement(signal_group, 'signal', {'refId': sig.get('id')})
    return signal_groups

def make_signal_control(signal_systems, signal_groups):
    ''' make signal control xml root '''
    signal_control = etree.Element('signalControl')
    for sig_sys in signal_systems.xpath('/signalSystems/signalSystem'):
        sig_sys_id = sig_sys.get('id')
        signal_system = make_subelement(signal_control, 'signalSystem', {'refId': sig_sys_id})
        sig_sys_controller = make_subelement(signal_system, 'signalSystemController')
        controller_identifier = make_subelement(sig_sys_controller, 'controllerIdentifier')
        controller_identifier.text = 'DefaultPlanbasedSignalSystemController'
        signal_plan = make_subelement(sig_sys_controller, 'signalPlan', {'id': '1'})
        make_subelement(signal_plan, 'start', {'daytime': '06:00:00'})
        make_subelement(signal_plan, 'stop', {'daytime': '18:00:00'})
        make_subelement(signal_plan, 'cycleTime', {'sec': str(CYCLE_TIME)})
        make_subelement(signal_plan, 'offset', {'sec': '0'})
        sig_groups = signal_groups.xpath("/signalGroups/signalSystem[@refId='{}']/signalGroup".format(sig_sys_id))
        for n_sig_grp, sig_grp in enumerate(sig_groups):
            sig_group_settings = make_subelement(signal_plan, 'signalGroupSettings', {'refId': sig_grp.get('id')})
            delta_sec = int(CYCLE_TIME / len(sig_groups))
            onset_sec = delta_sec * n_sig_grp
            dropping_sec = onset_sec + delta_sec - 5
            make_subelement(sig_group_settings, 'onset', {'sec': str(onset_sec)})
            make_subelement(sig_group_settings, 'dropping', {'sec': str(dropping_sec)})
    return signal_control

if __name__ == '__main__':
    # network nodes
    NETWORK_ROOT = etree.parse(INPUT_NETWORK)

    # signals nodes
    SIGNALS_ROOT = etree.parse(INPUT_SIGNALS)
    SIGNALS = SIGNALS_ROOT.xpath('/osm/node')

    # xml trees
    SIGNAL_SYSTEMS = make_signal_systems(NETWORK_ROOT, SIGNALS)
    SIGNAL_GROUPS = make_signal_groups(SIGNAL_SYSTEMS)
    SIGNAL_CONTROL = make_signal_control(SIGNAL_SYSTEMS, SIGNAL_GROUPS)

    # output
    TREE_SIGNAL_SYSTEMS = etree.ElementTree(SIGNAL_SYSTEMS)
    TREE_SIGNAL_GROUPS = etree.ElementTree(SIGNAL_GROUPS)
    TREE_SIGNAL_CONTROL = etree.ElementTree(SIGNAL_CONTROL)
    TREE_SIGNAL_SYSTEMS.write(OUT_SIGNAL_SYSTEMS, pretty_print=True, xml_declaration=True, encoding="utf-8")
    TREE_SIGNAL_GROUPS.write(OUT_SIGNAL_GROUPS, pretty_print=True, xml_declaration=True, encoding="utf-8")
    TREE_SIGNAL_CONTROL.write(OUT_SIGNAL_CONTROL, pretty_print=True, xml_declaration=True, encoding="utf-8")