diff --git a/flow/visualize/time_space_diagram.py b/flow/visualize/time_space_diagram.py index bc26ad855..3c7ab8b21 100644 --- a/flow/visualize/time_space_diagram.py +++ b/flow/visualize/time_space_diagram.py @@ -17,7 +17,7 @@ python time_space_diagram.py .csv .json """ from flow.utils.rllib import get_flow_params -from flow.networks import RingNetwork, FigureEightNetwork, MergeNetwork, I210SubNetwork +from flow.networks import RingNetwork, FigureEightNetwork, MergeNetwork, I210SubNetwork, HighwayNetwork import argparse from collections import defaultdict @@ -38,7 +38,8 @@ RingNetwork, FigureEightNetwork, MergeNetwork, - I210SubNetwork + I210SubNetwork, + HighwayNetwork ] @@ -103,7 +104,7 @@ def get_time_space_data(data, params): Returns ------- - ndarray (or dict of ndarray) + ndarray (or dict < str, np.ndarray >) 3d array (n_segments x 2 x 2) containing segments to be plotted. every inner 2d array is comprised of two 1d arrays representing [start time, start distance] and [end time, end distance] pairs. @@ -126,7 +127,8 @@ def get_time_space_data(data, params): RingNetwork: _ring_road, MergeNetwork: _merge, FigureEightNetwork: _figure_eight, - I210SubNetwork: _i210_subnetwork + I210SubNetwork: _i210_subnetwork, + HighwayNetwork: _highway, } # Get the function from switcher dictionary @@ -167,6 +169,33 @@ def _merge(data): return segs, data +def _highway(data): + r"""Generate time and position data for the highway. + + We generate plots for all lanes, so the segments are wrapped in + a dictionary. + + Parameters + ---------- + data : pd.DataFrame + cleaned dataframe of the trajectory data + + Returns + ------- + ndarray + 3d array (n_segments x 2 x 2) containing segments to be plotted. + every inner 2d array is comprised of two 1d arrays representing + [start time, start distance] and [end time, end distance] pairs. + pd.DataFrame + modified trajectory dataframe + """ + data.loc[:, :] = data[(data['distance'] > 500)] + data.loc[:, :] = data[(data['distance'] < 2300)] + segs = data[['time_step', 'distance', 'next_time', 'next_pos']].values.reshape((len(data), 2, 2)) + + return segs, data + + def _ring_road(data): r"""Generate position and speed data for the ring road. @@ -205,7 +234,7 @@ def _i210_subnetwork(data): Returns ------- - dict of ndarray + dict < str, np.ndarray > dictionary of 3d array (n_segments x 2 x 2) containing segments to be plotted. the dictionary is keyed on lane numbers, with the values being the 3d array representing the segments. every inner @@ -329,6 +358,8 @@ def _get_abs_pos(df, params): 'bottom_to_top': intersection / 2 + inner, 'right_to_left': junction + 3 * inner, } + elif params['network'] == HighwayNetwork: + return df['x'] else: edgestarts = defaultdict(float)