matsim_plan_gen/plan_gen.py

#!/usr/bin/env python3
''' Python generator for MATSim plans. '''

import sys
import numpy as np
import lxml.etree as etree

def read_nodes(input_network):
    ''' returns all network nodes as a list '''
    tree = etree.parse(input_network)
    return [node for node in tree.xpath("/network/nodes/node")]

def rand_person(nodes):
    ''' returns a person as a dict of random home and work coordinates '''
    len_nodes = len(nodes)
    home_node = nodes[np.random.randint(len_nodes)]
    work_node = nodes[np.random.randint(len_nodes)]
    home_xy = (home_node.get('x'), home_node.get('y'))
    work_xy = (work_node.get('x'), work_node.get('y'))
    return {'home': home_xy, 'work': work_xy}

def create_child(parent_node, child_name, child_attrs={}):
    ''' creates an xml child element and set its attributes '''
    child = etree.SubElement(parent_node, child_name)
    for attr, value in child_attrs.items():
        child.set(attr, value)
    return child

def make_plans(persons):
    ''' makes xml tree of plans based on persons list '''
    plans = etree.Element('plans')
    for n, p in enumerate(persons):
        person = create_child(plans, 'person', {'id': str(n+1)})
        plan = create_child(person, 'plan')
        # plan
        create_child(plan, 'act', {'type': 'h', 'x': p['home'][0], 'y': p['home'][1], 'end_time': '09:00:00'})
        create_child(plan, 'leg', {'mode': 'car'})
        create_child(plan, 'act', {'type': 'w', 'x': p['work'][0], 'y': p['work'][1], 'dur': '03:00:00'})
        create_child(plan, 'leg', {'mode': 'car'})
        create_child(plan, 'act', {'type': 'h', 'x': p['home'][0], 'y': p['home'][1]})
    return plans

if __name__ == '__main__':

    # command line arguments
    if len(sys.argv) != 3:
        print('usage:', sys.argv[0], '<input_network> <nb_persons>')
        sys.exit(-1)
    INPUT_NETWORK = sys.argv[1]
    NB_PERSONS = int(sys.argv[2])

    NODES = read_nodes(INPUT_NETWORK)
    PERSONS = [rand_person(NODES) for _ in range(NB_PERSONS)]
    PLANS = make_plans(PERSONS)

    # print XML
    print('<?xml version="1.0" ?>')
    print('<!DOCTYPE plans SYSTEM "http://www.matsim.org/files/dtd/plans_v4.dtd">')
    print(etree.tostring(PLANS, pretty_print=True).decode('utf-8'))