random plans (fixed hours)

plan_gen.py

@@ -1,22 +1,17 @@
 #!/usr/bin/env python3
-# -*- coding: utf-8 -*-
 ''' Python generator for MATSim plans. '''
 
+import sys
 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():
+def read_nodes(input_network):
     ''' returns all network nodes as a list '''
-    tree = etree.parse(INPUT_NETWORK)
+    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 '''
+    ''' 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)]
@@ -24,7 +19,41 @@ def rand_person(nodes):
     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__':
-    NODES = get_nodes()
+
+    # 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)]
-    print(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'))