--- a/python/storage.py	Wed Apr 15 18:17:50 2015 +0200
+++ b/python/storage.py	Thu Apr 16 11:34:51 2015 +0200
@@ -639,6 +639,9 @@
             finally: self.sechead = None
         else: return self.fp.readline()
 
+def generatePDLaneColumn(data):
+    data['LANE'] = data['LANE\LINK\NO'].astype(str)+'_'+data['LANE\INDEX'].astype(str)
+
 def loadTrajectoriesFromVissimFile(filename, simulationStepsPerTimeUnit, nObjects = -1, warmUpLastInstant = None, usePandas = False, nDecimals = 2):
     '''Reads data from VISSIM .fzp trajectory file
     simulationStepsPerTimeUnit is the number of simulation steps per unit of time used by VISSIM
@@ -653,7 +656,7 @@
         from pandas import read_csv
         from numpy import min, max, round
         data = read_csv(filename, delimiter=';', comment='*', header=0, skiprows = 1)
-        data['LANE'] = data['LANE\LINK\NO'].astype(str)+'_'+data['LANE\INDEX'].astype(str)
+        generatePDLaneColumn(data)
         data['TIME'] = data['$VEHICLE:SIMSEC']*simulationStepsPerTimeUnit
         grouped = data.loc[:,['NO','TIME']].groupby(['NO'], as_index = False)
         instants = grouped['TIME'].agg({'first': min, 'last': max})
@@ -692,18 +695,36 @@
 
         return objects.values()
 
-def countStoppedVehiclesVissim(filename, proportionStationaryTime = 0.7):
+def selectPDLanes(data, lanes = None):
+    '''Selects the subset of data for the right lanes
+
+    Lane format is a string 'x_y' where x is link index and y is lane index'''
+    if lanes is not None:
+        if 'LANE' not in data.columns:
+            generatePDLaneColumn(data)
+        indices = (data['LANE'] == lanes[0])
+        for l in lanes[1:]:
+            indices = indices | (data['LANE'] == l)
+        return data[indices]
+    else:
+        return data
+
+def countStoppedVehiclesVissim(filename, lanes = None, proportionStationaryTime = 0.7):
     '''Counts the number of vehicles stopped for a long time in a VISSIM trajectory file
     and the total number of vehicles
 
     Vehicles are considered finally stationary
-    if more than proportionStationaryTime of their total time'''
+    if more than proportionStationaryTime of their total time
+    If lanes is not None, only the data for the selected lanes will be provided
+    (format as string x_y where x is link index and y is lane index)'''
     from pandas import read_csv
     from numpy import array, sum as npsum
-    data = read_csv(filename, delimiter=';', comment='*', header=0, skiprows = 1, usecols = ['NO', '$VEHICLE:SIMSEC', 'POS']) # 'LANE\LINK\NO', 'LANE\INDEX', 
+    columns = ['NO', '$VEHICLE:SIMSEC', 'POS']
+    if lanes is not None:
+        columns += ['LANE\LINK\NO', 'LANE\INDEX']
+    data = read_csv(filename, delimiter=';', comment='*', header=0, skiprows = 1, usecols = columns)
+    data = selectPDLanes(data, lanes)
     data.sort(['$VEHICLE:SIMSEC'], inplace = True)
-    #merged = merge(data, data, how='inner', left_on=['LANE\LINK\NO', 'LANE\INDEX', '$VEHICLE:SIMSEC'], right_on=['LANE\LINK\NO', 'LANE\INDEX', '$VEHICLE:SIMSEC'], sort = False)
-    #merged = merged[merged['NO_x']>merged['NO_y']]
 
     nStationary = 0
     from matplotlib.pyplot import plot, figure
@@ -717,14 +738,17 @@
 
     return nStationary, nVehicles
 
-def countCollisionsVissim(filename, collisionTimeDifference = 0.2):
+def countCollisionsVissim(filename, lanes = None, collisionTimeDifference = 0.2):
     '''Counts the number of collisions per lane in a VISSIM trajectory file
 
     To distinguish between cars passing and collision, 
     one checks when the sign of the position difference inverts
-    (if the time are closer than collisionTimeDifference)'''
+    (if the time are closer than collisionTimeDifference)
+    If lanes is not None, only the data for the selected lanes will be provided
+    (format as string x_y where x is link index and y is lane index)'''
     from pandas import read_csv, merge
     data = read_csv(filename, delimiter=';', comment='*', header=0, skiprows = 1, usecols = ['LANE\LINK\NO', 'LANE\INDEX', '$VEHICLE:SIMSEC', 'NO', 'POS'])
+    data = selectPDLanes(data, lanes)
     merged = merge(data, data, how='inner', left_on=['LANE\LINK\NO', 'LANE\INDEX', '$VEHICLE:SIMSEC'], right_on=['LANE\LINK\NO', 'LANE\INDEX', '$VEHICLE:SIMSEC'], sort = False)
     merged = merged[merged['NO_x']>merged['NO_y']]