sumo.py

#! /usr/bin/env python
'''Libraries for the SUMO traffic simulation software
http://sumo.dlr.de
'''
import pandas as pd
import sys
#import csv

''' This is the valid one'''


def loadTazEdges(inFilename):
	'''Converts list of SUMO edges per OSM id and groups per TAZ id
		The input must be a csv file with first two columns the OSM id and TAZ id, then the list of SUMO edge id
		Returns a dictionary with keys = TAZ id and values = list of SUMO edges per TAZ 
		inFilename = taz_edges.csv '''
	data = []
	tazs = {}
	with open(inFilename,'r') as f:
		f.readline() # skip the headers
		for r in f:
			tmp = r.strip().split(',')
			tazID = tmp[1]
			for edge in tmp[2:]:                
				if len(edge) > 0:
					if tazID in tazs:
						if edge not in tazs[tazID]:
							tazs[tazID].append(edge)
					else:
						tazs[tazID] = [edge]
	f.close()
	return tazs

def edge2Taz(tazs):
	'''Returns the associative array of the TAZ of each SUMO edge
		It is a dictionary with keys = SUMO edge id and values = TAZ id
		The input must be a file generated by the loadTazEdges method'''
	edge2Tazs = {}
	for taz, edges in tazs.iteritems():
		for edge in edges:
			if edge in edge2Tazs:
				print('error for edge: {} (taz {}/{})'.format(edge, edge2Tazs[edge], taz))
			edge2Tazs[edge] = taz
	return edge2Tazs

def saveTazEdges(outFilename, tazs):
	'''Produces a xml formatted file to be load in SUMO as taz file 
		The input must be a file generated by the loadTazEdges method'''
	with open(outFilename,'w') as out:
		out.write('<tazs>\n')				
		for tazID in tazs:
			out.write('<taz id="{}" edges="'.format(tazID)+' '.join(tazs[tazID])+'"/>\n')
		out.write('</tazs>\n')
	out.close()

def tripsPerTaz(tripFile):
	'''Returns the number of trips connecting the different TAZ
		The input must be a tripinfo csv formatted file and the taz file used in the simulation
		Returns a dataframe containing the trips count from TAZ to TAZ
		tripFile = tripInfo1.csv '''
	tazs = loadTazEdges('taz_edges.csv')
	edgeTAZ = edge2Taz(tazs)
	trips_tazs = {}
	trips_edges = {}
	print 'reading tripinfo file'
	df = pd.read_csv(tripFile)
	trips = df[['tripinfo_id', 'tripinfo_departLane', 'tripinfo_arrivalLane']]
	#tripsLanes = trips.set_index('tripinfo_id')
	for i in range(len(trips)):
		tripID = trips.iloc[i]['tripinfo_id']
		trips_edges[tripID] =[trips.iloc[i]['tripinfo_departLane'], trips.iloc[i]['tripinfo_arrivalLane']]
	
	print 'trips_edges built'
	
	for trip in trips_edges.keys():
		fromEdge = trips_edges.get(trip)[0][:-2]
		fromTAZ = edgeTAZ.get(fromEdge)
		toEdge = trips_edges.get(trip)[1][:-2]
		toTAZ = edgeTAZ.get(toEdge)
		trips_tazs[trip] = [fromTAZ, toTAZ]
	df = pd.DataFrame(data=trips_tazs, index=['fromTAZ', 'toTAZ']).T
	df.to_csv('tripsTaz.csv')
	print df
		
def readTazFile(tazFile):
	'''Constructs a dictionary containing the TAZ id and the SUMO edges per TAZ id
       The input mus be the taz file used in the simulation'''
	taz_dict = {}
	tazIDList = []
	with open(tazFile,'r') as f:
		lines = f.readlines()
		f.close()
		taz_list = lines[1:-1]
		for line in taz_list:
			tmp = line.split(' ')
			tazID = tmp[1].split('"')[1]
			tazIDList.append(tazID)
			firstEdge = tmp[2].split('"')[1]
			if len(tmp) == 3:
				edges = [firstEdge]
			else:
				lastEdge = tmp[len(tmp)-1].split('"')[0]
				if len(tmp) == 4:
					edges = [firstEdge, lastEdge]
				else:
					edges = tmp[3:-1]
					edges.insert(0,firstEdge)
					edges.append(lastEdge)
			if taz_dict.has_key(tazID):
				print 'tazID', tazID, 'duplicated'
				sys.exit()
			else:
				taz_dict[tazID] = edges
	return taz_dict