diff --git a/flow/core/experiment.py b/flow/core/experiment.py index 464b0a405..a7ac07738 100755 --- a/flow/core/experiment.py +++ b/flow/core/experiment.py @@ -2,6 +2,7 @@ from flow.utils.registry import make_create_env from flow.data_pipeline.data_pipeline import write_dict_to_csv, upload_to_s3, get_extra_info, get_configuration from flow.data_pipeline.leaderboard_utils import network_name_translate +from flow.visualize.time_space_diagram import tsd_main from collections import defaultdict from datetime import datetime, timezone import logging @@ -20,8 +21,8 @@ class Experiment: the actions of RL agents in the network, type the following: >>> from flow.envs import Env - >>> flow_params = dict(...) # see the examples in exp_config - >>> exp = Experiment(flow_params) # for some experiment configuration + {'network': >>> self.env.network.__class__} = dict(...) # see the examples in exp_config + {'network': >>> exp = Experiment(self.env.network.__class__}) # for some experiment configuration >>> exp.run(num_runs=1) If you wish to specify the actions of RL agents in the network, this may be @@ -39,7 +40,7 @@ class can generate csv files from emission files produced by sumo. These ``emission_path`` attribute in ``SimParams`` to some path. >>> from flow.core.params import SimParams - >>> flow_params['sim'] = SimParams(emission_path="./data") + {'network': >>> self.env.network.__class__}['sim'] = SimParams(emission_path="./data") Once you have included this in your environment, run your Experiment object as follows: @@ -233,6 +234,11 @@ def rl_actions(*_): write_dict_to_csv(metadata_table_path, metadata, True) if to_aws: + tsd_main(trajectory_table_path, + {'network': self.env.network.__class__}, + min_speed=0, + max_speed=10, + start=self.env.env_params.warmup_steps) upload_to_s3('circles.data.pipeline', 'metadata_table/date={0}/partition_name={1}_METADATA/{1}_METADATA.csv'.format(cur_date, source_id), @@ -241,5 +247,8 @@ def rl_actions(*_): 'fact_vehicle_trace/date={0}/partition_name={1}/{1}.csv'.format(cur_date, source_id), trajectory_table_path, {'network': metadata['network'][0], 'is_baseline': metadata['is_baseline'][0]}) + upload_to_s3('circles.data.pipeline', + 'time_space_diagram/date={0}/partition_name={1}/{1}.png'.format(cur_date, source_id), + trajectory_table_path.replace('csv', 'png')) return info_dict diff --git a/flow/visualize/time_space_diagram.py b/flow/visualize/time_space_diagram.py index b1500b48d..a9392e21d 100644 --- a/flow/visualize/time_space_diagram.py +++ b/flow/visualize/time_space_diagram.py @@ -382,7 +382,7 @@ def _get_abs_pos(df, params): return ret -def plot_tsd(ax, df, segs, args, lane=None, ghost_edges=None, ghost_bounds=None): +def plot_tsd(ax, df, segs, cmap, min_speed=0, max_speed=10, start=0, lane=None, ghost_edges=None, ghost_bounds=None): """Plot the time-space diagram. Take the pre-processed segments and other meta-data, then plot all the line segments. @@ -395,8 +395,12 @@ def plot_tsd(ax, df, segs, args, lane=None, ghost_edges=None, ghost_bounds=None) data used for axes bounds and speed coloring segs : list of list of lists line segments to be plotted, where each segment is a list of two [x,y] pairs - args : dict - parsed arguments + min_speed : int or float + minimum speed in colorbar + max_speed : int or float + maximum speed in colorbar + start : int or float + starting time_step not greyed out lane : int, optional lane number to be shown in plot title ghost_edges : list or set of str @@ -408,7 +412,7 @@ def plot_tsd(ax, df, segs, args, lane=None, ghost_edges=None, ghost_bounds=None) ------- None """ - norm = plt.Normalize(args.min_speed, args.max_speed) + norm = plt.Normalize(min_speed, max_speed) xmin, xmax = df['time_step'].min(), df['time_step'].max() xbuffer = (xmax - xmin) * 0.025 # 2.5% of range @@ -418,7 +422,7 @@ def plot_tsd(ax, df, segs, args, lane=None, ghost_edges=None, ghost_bounds=None) ax.set_xlim(xmin - xbuffer, xmax + xbuffer) ax.set_ylim(ymin - ybuffer, ymax + ybuffer) - lc = LineCollection(segs, cmap=my_cmap, norm=norm) + lc = LineCollection(segs, cmap=cmap, norm=norm) lc.set_array(df['speed'].values) lc.set_linewidth(1) ax.add_collection(lc) @@ -428,15 +432,15 @@ def plot_tsd(ax, df, segs, args, lane=None, ghost_edges=None, ghost_bounds=None) if ghost_edges: y_domain_min = df[~df['edge_id'].isin(ghost_edges)]['distance'].min() y_domain_max = df[~df['edge_id'].isin(ghost_edges)]['distance'].max() - rects.append(Rectangle((xmin, y_domain_min), args.start - xmin, y_domain_max - y_domain_min)) + rects.append(Rectangle((xmin, y_domain_min), start - xmin, y_domain_max - y_domain_min)) rects.append(Rectangle((xmin, ymin), xmax - xmin, y_domain_min - ymin)) rects.append(Rectangle((xmin, y_domain_max), xmax - xmin, ymax - y_domain_max)) elif ghost_bounds: - rects.append(Rectangle((xmin, ghost_bounds[0]), args.start - xmin, ghost_bounds[1] - ghost_bounds[0])) + rects.append(Rectangle((xmin, ghost_bounds[0]), start - xmin, ghost_bounds[1] - ghost_bounds[0])) rects.append(Rectangle((xmin, ymin), xmax - xmin, ghost_bounds[0] - ymin)) rects.append(Rectangle((xmin, ghost_bounds[1]), xmax - xmin, ymax - ghost_bounds[1])) else: - rects.append(Rectangle((xmin, ymin), args.start - xmin, ymax - ymin)) + rects.append(Rectangle((xmin, ymin), start - xmin, ymax - ymin)) if rects: pc = PatchCollection(rects, facecolor='grey', alpha=0.5, edgecolor=None) @@ -457,41 +461,28 @@ def plot_tsd(ax, df, segs, args, lane=None, ghost_edges=None, ghost_bounds=None) cbar.ax.tick_params(labelsize=18) -if __name__ == '__main__': - # create the parser - parser = argparse.ArgumentParser( - formatter_class=argparse.RawDescriptionHelpFormatter, - description='[Flow] Generates time space diagrams for flow networks.', - epilog='python time_space_diagram.py .csv ' - '.json') - - # required arguments - parser.add_argument('trajectory_path', type=str, - help='path to the Flow trajectory csv file.') - parser.add_argument('flow_params', type=str, - help='path to the flow_params json file.') - - # optional arguments - parser.add_argument('--steps', type=int, default=1, - help='rate at which steps are plotted.') - parser.add_argument('--title', type=str, default='Time Space Diagram', - help='rate at which steps are plotted.') - parser.add_argument('--max_speed', type=int, default=8, - help='The maximum speed in the color range.') - parser.add_argument('--min_speed', type=int, default=0, - help='The minimum speed in the color range.') - parser.add_argument('--start', type=float, default=0, - help='initial time (in sec) in the plot.') - - args = parser.parse_args() - - # flow_params is imported as a dictionary - if '.json' in args.flow_params: - flow_params = get_flow_params(args.flow_params) - else: - module = __import__("examples.exp_configs.non_rl", fromlist=[args.flow_params]) - flow_params = getattr(module, args.flow_params).flow_params +def tsd_main(trajectory_path, flow_params, min_speed=0, max_speed=10, start=0): + """Prepare and plot the time-space diagram. + Parameters + ---------- + trajectory_path : str + file path (for the .csv formatted file) + flow_params : dict + flow-specific parameters, including: + * "network" (str): name of the network that was used when generating + the emission file. Must be one of the network names mentioned in + ACCEPTABLE_NETWORKS, + * "net_params" (flow.core.params.NetParams): network-specific + parameters. This is used to collect the lengths of various network + links. + min_speed : int or float + minimum speed in colorbar + max_speed : int or float + maximum speed in colorbar + start : int or float + starting time_step not greyed out + """ # some plotting parameters cdict = { 'red': ((0, 0, 0), (0.2, 1, 1), (0.6, 1, 1), (1, 0, 0)), @@ -501,29 +492,50 @@ def plot_tsd(ax, df, segs, args, lane=None, ghost_edges=None, ghost_bounds=None) my_cmap = colors.LinearSegmentedColormap('my_colormap', cdict, 1024) # Read trajectory csv into pandas dataframe - traj_df = import_data_from_trajectory(args.trajectory_path, flow_params) + traj_df = import_data_from_trajectory(trajectory_path, flow_params) # Convert df data into segments for plotting segs, traj_df = get_time_space_data(traj_df, flow_params) if flow_params['network'] == I210SubNetwork: nlanes = traj_df['lane_id'].nunique() - fig = plt.figure(figsize=(16, 9*nlanes)) + plt.figure(figsize=(16, 9*nlanes)) for lane, df in traj_df.groupby('lane_id'): ax = plt.subplot(nlanes, 1, lane+1) - plot_tsd(ax, df, segs[lane], args, int(lane+1), ghost_edges={'ghost0', '119257908#3'}) + plot_tsd(ax=ax, + df=df, + segs=segs[lane], + cmap=my_cmap, + min_speed=min_speed, + max_speed=max_speed, + start=start, + lane=int(lane+1), + ghost_edges={'ghost0', '119257908#3'}) plt.tight_layout() else: # perform plotting operation - fig = plt.figure(figsize=(16, 9)) + plt.figure(figsize=(16, 9)) ax = plt.axes() if flow_params['network'] == HighwayNetwork: - plot_tsd(ax, traj_df, segs, args, ghost_bounds=(500, 2300)) + plot_tsd(ax=ax, + df=traj_df, + segs=segs, + cmap=my_cmap, + min_speed=min_speed, + max_speed=max_speed, + start=start, + ghost_bounds=(500, 2300)) else: - plot_tsd(ax, traj_df, segs, args) + plot_tsd(ax=ax, + df=traj_df, + segs=segs, + cmap=my_cmap, + min_speed=min_speed, + max_speed=max_speed, + start=start) ########################################################################### # Note: For MergeNetwork only # @@ -534,5 +546,43 @@ def plot_tsd(ax, df, segs, args, lane=None, ghost_edges=None, ghost_bounds=None) [-0.1, -0.1], linewidth=3, color="white") # ########################################################################### - outfile = args.trajectory_path.replace('csv', 'png') + outfile = trajectory_path.replace('csv', 'png') plt.savefig(outfile) + + +if __name__ == '__main__': + # create the parser + parser = argparse.ArgumentParser( + formatter_class=argparse.RawDescriptionHelpFormatter, + description='[Flow] Generates time space diagrams for flow networks.', + epilog='python time_space_diagram.py .csv ' + '.json') + + # required arguments + parser.add_argument('trajectory_path', type=str, + help='path to the Flow trajectory csv file.') + parser.add_argument('flow_params', type=str, + help='path to the flow_params json file.') + + # optional arguments + parser.add_argument('--steps', type=int, default=1, + help='rate at which steps are plotted.') + parser.add_argument('--title', type=str, default='Time Space Diagram', + help='rate at which steps are plotted.') + parser.add_argument('--max_speed', type=int, default=8, + help='The maximum speed in the color range.') + parser.add_argument('--min_speed', type=int, default=0, + help='The minimum speed in the color range.') + parser.add_argument('--start', type=float, default=0, + help='initial time (in sec) in the plot.') + + args = parser.parse_args() + + # flow_params is imported as a dictionary + if '.json' in args.flow_params: + flow_params = get_flow_params(args.flow_params) + else: + module = __import__("examples.exp_configs.non_rl", fromlist=[args.flow_params]) + flow_params = getattr(module, args.flow_params).flow_params + + tsd_main(args.trajectory_path, flow_params, min_speed=args.min_speed, max_speed=args.max_speed, start=args.start)