matsim_plan_gen/plan_gen.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
''' Python generator for MATSim plans. '''

import numpy as np
import lxml.etree as etree

# constants
INPUT_NETWORK = 'input/network.xml'
OUTPUT_PLANS = 'out_plans.xml'
NB_PERSONS = 10

def get_nodes():
    ''' 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 home and work random locations '''
    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}

if __name__ == '__main__':
    NODES = get_nodes()
    PERSONS = [rand_person(NODES) for _ in range(NB_PERSONS)]
    print(PERSONS)