clusters (not tested)

plan_gen.py

@@ -13,6 +13,8 @@ WORK_DURATION = '04:00:00'
 
 def read_nodes(input_network):
     ''' returns all network nodes as a list '''
+    if not input_network:
+        return None
     tree = etree.parse(input_network)
     return [node for node in tree.xpath("/network/nodes/node")]
 
@@ -21,16 +23,19 @@ def rand_time(low, high):
     delta = np.random.randint(high - low)
     return time.strftime('%H:%M:%S', time.gmtime(low + delta))
 
-def rand_person(nodes):
+def rand_node_xy(nodes, clusters):
+    ''' returns a random node coordinates from a list of nodes '''
+    used_nodes = nodes
+    if any(clusters):
+        cluster = np.random.randint(len(clusters))
+        used_nodes = clusters[cluster]
+    node = used_nodes[np.random.randint(len(used_nodes))]
+    return (node.get('x'), node.get('y'))
+
+def rand_person(nodes, home_clusters, work_clusters):
     ''' returns a person as a dictionnary of random parameters '''
-    len_nodes = len(nodes)
-    # home coordinates
-    home_node = nodes[np.random.randint(len_nodes)]
-    home_xy = (home_node.get('x'), home_node.get('y'))
-    # work coordinates
-    work_node = nodes[np.random.randint(len_nodes)]
-    work_xy = (work_node.get('x'), work_node.get('y'))
-    # home departure time
+    home_xy = rand_node_xy(nodes, home_clusters)
+    work_xy = rand_node_xy(nodes, work_clusters)
     home_departure = rand_time(MIN_DEPARTURE_TIME, MAX_DEPARTURE_TIME)
     return {'home': home_xy, 'work': work_xy, 'home_departure': home_departure}
 
@@ -60,14 +65,22 @@ def make_plans(persons):
 if __name__ == '__main__':
 
     # command line arguments
-    if len(sys.argv) != 3:
-        print('usage:', sys.argv[0], '<input_network> <nb_persons>')
+    if len(sys.argv) != 5:
+        print('usage:', sys.argv[0], '<nb_persons> <input_network> <input_home_clusters> <input_work_clusters>')
+        print('type "" if you don\'t want to use specific clusters')
         sys.exit(-1)
-    INPUT_NETWORK = sys.argv[1]
-    NB_PERSONS = int(sys.argv[2])
+    NB_PERSONS = int(sys.argv[1])
+    INPUT_NETWORK = sys.argv[2]
+    INPUT_HOME_CLUSTERS = sys.argv[3].split(',')
+    INPUT_WORK_CLUSTERS = sys.argv[4].split(',')
 
+    # get data
     NODES = read_nodes(INPUT_NETWORK)
-    PERSONS = [rand_person(NODES) for _ in range(NB_PERSONS)]
+    HOME_CLUSTERS = [read_nodes(cluster) for cluster in INPUT_HOME_CLUSTERS]
+    WORK_CLUSTERS = [read_nodes(cluster) for cluster in INPUT_WORK_CLUSTERS]
+
+    # make xml
+    PERSONS = [rand_person(NODES, HOME_CLUSTERS, WORK_CLUSTERS) for _ in range(NB_PERSONS)]
     PLANS = make_plans(PERSONS)
 
     # print XML