relay-throughput-testing/src/plot_streams.py

import gzip
import pickle
import matplotlib.pylab as plt
import numpy as np
#
def plot_cells(stream, plot_widget, time_offset=None, clickable=False, label=None, color=None):
	num_bytes = np.cumsum(stream['length'])
	timestamps = np.asarray(stream['timestamp'])
	#
	if time_offset is not None:
		timestamps = timestamps-time_offset
	#
	return ax.step(timestamps, num_bytes/(1024**2), where='post', label=label, color=color, picker=(5 if clickable else None))[0]
#
def onresize(event):
	# matplotlib axes size only scales based on fractions (even with tight_layout), so we manually calculate a fixed padding size
	w = event.width/event.canvas.figure.dpi
	h = event.height/event.canvas.figure.dpi
	w_padding = 0.8 # in inches
	h_padding = 0.6 # in inches
	#
	for ax in event.canvas.figure.axes:
		for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] + ax.get_xticklabels() + ax.get_yticklabels()):
			item.set_fontsize(11+0.5*(w*h/(5**2)))
		#
	#
	event.canvas.figure.subplots_adjust(left=w_padding/w+0.01, right=1-((w_padding/2)/w), top=1-(h_padding/h), bottom=h_padding/h+0.01)
#
def onscroll(event):
	scale = 2**(-event.step)
	ax = event.inaxes
	mouse_x = event.xdata
	mouse_y = event.ydata
	old_xlim = ax.get_xlim()
	old_ylim = ax.get_ylim()
	#
	ax.set_xlim([old_xlim[0]*scale+mouse_x*(1-scale), old_xlim[1]*scale+mouse_x*(1-scale)])
	ax.set_ylim([old_ylim[0]*scale+mouse_y*(1-scale), old_ylim[1]*scale+mouse_y*(1-scale)])
	#
	event.canvas.draw()
#
def onpick(event, lines):
	this_line = event.artist
	if event.mouseevent.button == 1 and not event.mouseevent.dblclick and event.mouseevent.key is None:
		# if the mouse is single-clicked and no keyboard key is held down
		if this_line.get_visible() and this_line in lines.keys():
			this_info = [x[1] for x in lines.items() if x[0] == this_line][0]
			#
			for_legend = []
			for (line, info) in lines.items():
				if info['measureme_id'] != this_info['measureme_id']:
					line.set_visible(False)
				else:
					line.set_visible(True)
					for_legend.append(line)
				#
			#
			event.mouseevent.inaxes.legend(for_legend, [x.get_label() for x in for_legend], loc='upper right')
			event.canvas.draw_idle()
		#
	#
#
def onclick(event, lines):
	if event.button == 1 and event.dblclick and event.key is None:
		# if the mouse is double-clicked and no keyboard key is held down
		for (line, info) in lines.items():
			if info['is_main_plot']:
				line.set_visible(True)
			else:
				line.set_visible(False)
			#
		#
		event.inaxes.get_legend().remove()
		event.canvas.draw_idle()
	#
#
def onmotion(event, pan_settings):
	if event.inaxes is not None and event.key == 'control' and event.button == 1:
		ax = event.inaxes
		pixel_to_data = ax.transData.inverted()
		current_pos = pixel_to_data.transform_point((event.x, event.y))
		last_pos = pixel_to_data.transform_point((pan_settings['start_x'], pan_settings['start_y']))
		#
		old_xlim = ax.get_xlim()
		old_ylim = ax.get_ylim()
		#
		ax.set_xlim([old_xlim[0]+(last_pos[0]-current_pos[0]), old_xlim[1]+(last_pos[0]-current_pos[0])])
		ax.set_ylim([old_ylim[0]+(last_pos[1]-current_pos[1]), old_ylim[1]+(last_pos[1]-current_pos[1])])
		#
		event.canvas.draw_idle()
	#
	pan_settings['start_x'] = event.x
	pan_settings['start_y'] = event.y
#
def onkeypress(event, lines):
	if event.key == 'control':
		num_points = 0
		range_x = event.inaxes.xaxis.get_view_interval()
		range_y = event.inaxes.yaxis.get_view_interval()
		#
		for line in [l for l in lines if l.get_visible()]:
			data_x = line.get_xdata(orig=True)
			data_y = line.get_ydata(orig=True)
			#
			num_points += ((data_x>=range_x[0]) & (data_x<=range_x[1]) & (data_y>=range_y[0]) & (data_y<=range_y[1])).sum()
			# how many points are being rendered
		#
		for (line, info) in lines.items():
			data_x = line.get_xdata(orig=True)
			data_y = line.get_ydata(orig=True)
			line.orig_x = data_x
			line.orig_y = data_y
			line.orig_drawstyle = line.get_drawstyle()
			#
			subsample_spacing = max(1, int(num_points/10000))
			# the constant can be decreased to speed up plotting on slower computers
			#
			mask = np.ones(len(data_x))
			mask[::subsample_spacing] = 0
			line.set_xdata(data_x[mask==0])
			line.set_ydata(data_y[mask==0])
			line.set_drawstyle('default')
		#
		event.canvas.draw_idle()
	#
#
def onkeyrelease(event, lines):
	if event.key == 'control':
		for (line, info) in lines.items():
			line.set_xdata(line.orig_x)
			line.set_ydata(line.orig_y)
			line.set_drawstyle(line.orig_drawstyle)
		#
		event.canvas.draw_idle()
	#
#
def get_complimentary_color(color_index):
	return (color_index+1 if color_index%2==0 else color_index-1)
#
if __name__ == '__main__':
	with gzip.GzipFile('processed-data.pickle.gz', 'rb') as f:
		streams = pickle.load(f)
	#
	fig, ax = plt.subplots()#constrained_layout=True
	#
	start_time = min([hop[t]['timestamp'][0] for m in streams for s in streams[m] for d in streams[m][s]
	                  for hop in streams[m][s][d] for t in ('received','sent')])
	#
	lines = {}
	assigned_colors = []
	colormap = plt.get_cmap('tab20') #'tab10'
	direction_shortforms = {'forward':'fwd', 'backward':'bwd'}
	transmission_shortforms = {'received':'recv', 'sent':'sent'}
	#
	for measureme_id in streams:
		# for each circuit
		for stream_id in streams[measureme_id]:
			# for each stream
			direction = 'forward'
			for hop_index in range(len(streams[measureme_id][stream_id][direction])):
				# for each hop in the circuit (including the OP)
				for transmission in ('received', 'sent'):
					data = streams[measureme_id][stream_id][direction][hop_index][transmission]
					#
					if hop_index == len(streams[measureme_id][stream_id][direction])-1:
						guard_fingerprint = streams[measureme_id][stream_id][direction][1]['fingerprint']
						if guard_fingerprint not in assigned_colors:
							assigned_colors.append(guard_fingerprint)
						#
						if transmission == 'sent':
							color_index = assigned_colors.index(guard_fingerprint)
						else:
							color_index = get_complimentary_color(assigned_colors.index(guard_fingerprint))
						#
					else:
						color_index = hop_index*2 + ('received', 'sent').index(transmission)
					#
					is_main_plot = (hop_index == len(streams[measureme_id][stream_id][direction])-1 and transmission == 'sent')
					direction_label = direction_shortforms[direction]
					transmission_label = transmission_shortforms[transmission]
					label = 'hop={}, {:.4}, {:.4}, mid={}, sid={}'.format(hop_index, direction_label, transmission_label, measureme_id, stream_id)
					line = plot_cells(data, ax, time_offset=start_time, clickable=is_main_plot, label=label, color=colormap(color_index))
					if not is_main_plot:
						line.set_visible(False)
					#
					lines[line] = {'measureme_id':measureme_id, 'stream_id':stream_id, 'direction':direction,
					               'hop_index':hop_index, 'transmission':transmission, 'is_main_plot':is_main_plot}
				#
			#
		#
	#
	ax.set_xlabel('Time (s)')
	ax.set_ylabel('Data (MiB)')
	ax.set_title('Test')
	fig.tight_layout(pad=0)
	#ax.set_ylim(0, None)
	#
	fig.canvas.mpl_connect('resize_event', onresize)
	fig.canvas.mpl_connect('scroll_event', onscroll)
	fig.canvas.mpl_connect('pick_event', lambda event,lines=lines: onpick(event, lines))
	fig.canvas.mpl_connect('button_press_event', lambda event,lines=lines: onclick(event, lines))
	fig.canvas.mpl_connect('motion_notify_event', lambda event,pan_settings={'start_x':0,'start_y':0}: onmotion(event, pan_settings))
	fig.canvas.mpl_connect('key_press_event', lambda event,lines=lines: onkeypress(event, lines))
	fig.canvas.mpl_connect('key_release_event', lambda event,lines=lines: onkeyrelease(event, lines))
	#
	plt.show(fig)
#