1#!/usr/bin/env python3.4 2# 3# Copyright 2021 - The Android Open Source Project 4# 5# Licensed under the Apache License, Version 2.0 (the 'License'); 6# you may not use this file except in compliance with the License. 7# You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, software 12# distributed under the License is distributed on an 'AS IS' BASIS, 13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14# See the License for the specific language governing permissions and 15# limitations under the License. 16 17import bokeh, bokeh.plotting, bokeh.io 18import collections 19import itertools 20import json 21import math 22import numpy 23from acts_contrib.test_utils.wifi import wifi_performance_test_utils as wputils 24 25 26# Plotting Utilities 27class BokehFigure(): 28 """Class enabling simplified Bokeh plotting.""" 29 30 COLORS = [ 31 'black', 32 'blue', 33 'blueviolet', 34 'brown', 35 'burlywood', 36 'cadetblue', 37 'cornflowerblue', 38 'crimson', 39 'cyan', 40 'darkblue', 41 'darkgreen', 42 'darkmagenta', 43 'darkorange', 44 'darkred', 45 'deepskyblue', 46 'goldenrod', 47 'green', 48 'grey', 49 'indigo', 50 'navy', 51 'olive', 52 'orange', 53 'red', 54 'salmon', 55 'teal', 56 'yellow', 57 ] 58 MARKERS = [ 59 'asterisk', 'circle', 'circle_cross', 'circle_x', 'cross', 'diamond', 60 'diamond_cross', 'hex', 'inverted_triangle', 'square', 'square_x', 61 'square_cross', 'triangle', 'x' 62 ] 63 64 TOOLS = ('box_zoom,box_select,pan,crosshair,redo,undo,reset,hover,save') 65 66 def __init__(self, 67 title=None, 68 x_label=None, 69 primary_y_label=None, 70 secondary_y_label=None, 71 height=700, 72 width=1100, 73 title_size='15pt', 74 axis_label_size='12pt', 75 legend_label_size='12pt', 76 legend_location = 'top_right', 77 axis_tick_label_size='12pt', 78 x_axis_type='auto', 79 sizing_mode='scale_both', 80 json_file=None): 81 if json_file: 82 self.load_from_json(json_file) 83 else: 84 self.figure_data = [] 85 self.fig_property = { 86 'title': title, 87 'x_label': x_label, 88 'primary_y_label': primary_y_label, 89 'secondary_y_label': secondary_y_label, 90 'num_lines': 0, 91 'height': height, 92 'width': width, 93 'title_size': title_size, 94 'axis_label_size': axis_label_size, 95 'legend_label_size': legend_label_size, 96 'legend_location': legend_location, 97 'axis_tick_label_size': axis_tick_label_size, 98 'x_axis_type': x_axis_type, 99 'sizing_mode': sizing_mode 100 } 101 102 def init_plot(self): 103 self.plot = bokeh.plotting.figure( 104 sizing_mode=self.fig_property['sizing_mode'], 105 width=self.fig_property['width'], 106 height=self.fig_property['height'], 107 title=self.fig_property['title'], 108 tools=self.TOOLS, 109 x_axis_type=self.fig_property['x_axis_type'], 110 output_backend='webgl') 111 tooltips = [ 112 ('index', '$index'), 113 ('(x,y)', '($x, $y)'), 114 ] 115 hover_set = [] 116 for line in self.figure_data: 117 hover_set.extend(line['hover_text'].keys()) 118 hover_set = set(hover_set) 119 for item in hover_set: 120 tooltips.append((item, '@{}'.format(item))) 121 self.plot.hover.tooltips = tooltips 122 self.plot.add_tools( 123 bokeh.models.tools.WheelZoomTool(dimensions='width')) 124 self.plot.add_tools( 125 bokeh.models.tools.WheelZoomTool(dimensions='height')) 126 127 def _filter_line(self, x_data, y_data, hover_text=None, error_bars=[]): 128 """Function to remove NaN points from bokeh plots.""" 129 x_data_filtered = [] 130 y_data_filtered = [] 131 error_bars_data_filtered = [] 132 hover_text_filtered = {} 133 for idx, xy in enumerate( 134 itertools.zip_longest(x_data, y_data, error_bars, fillvalue=float('nan'))): 135 if not math.isnan(xy[1]): 136 x_data_filtered.append(xy[0]) 137 y_data_filtered.append(xy[1]) 138 if len(error_bars)>0: 139 error_bars_data_filtered.append(xy[2]) 140 if hover_text: 141 for key, value in hover_text.items(): 142 hover_text_filtered.setdefault(key, []) 143 hover_text_filtered[key].append( 144 value[idx] if len(value) > idx else '') 145 return x_data_filtered, y_data_filtered, hover_text_filtered, error_bars_data_filtered 146 147 def add_line(self, 148 x_data, 149 y_data, 150 legend, 151 hover_text=None, 152 color=None, 153 width=3, 154 style='solid', 155 marker=None, 156 marker_size=10, 157 shaded_region=None, 158 error_bars=[], 159 y_axis='default'): 160 """Function to add line to existing BokehFigure. 161 162 Args: 163 x_data: list containing x-axis values for line 164 y_data: list containing y_axis values for line 165 legend: string containing line title 166 hover_text: text to display when hovering over lines 167 color: string describing line color 168 width: integer line width 169 style: string describing line style, e.g, solid or dashed 170 marker: string specifying line marker, e.g., cross 171 shaded region: data describing shaded region to plot 172 y_axis: identifier for y-axis to plot line against 173 """ 174 if y_axis not in ['default', 'secondary']: 175 raise ValueError('y_axis must be default or secondary') 176 if color == None: 177 color = self.COLORS[self.fig_property['num_lines'] % 178 len(self.COLORS)] 179 if style == 'dashed': 180 style = [5, 5] 181 if isinstance(hover_text, list): 182 hover_text = {'info': hover_text} 183 x_data_filter, y_data_filter, hover_text_filter, error_bars_filter = self._filter_line( 184 x_data, y_data, hover_text, error_bars) 185 self.figure_data.append({ 186 'x_data': x_data_filter, 187 'y_data': y_data_filter, 188 'legend': legend, 189 'hover_text': hover_text_filter, 190 'color': color, 191 'width': width, 192 'style': style, 193 'marker': marker, 194 'marker_size': marker_size, 195 'shaded_region': shaded_region, 196 'error_bars': error_bars_filter, 197 'y_axis': y_axis 198 }) 199 self.fig_property['num_lines'] += 1 200 201 def add_scatter(self, 202 x_data, 203 y_data, 204 legend, 205 hover_text=None, 206 color=None, 207 marker=None, 208 marker_size=10, 209 y_axis='default'): 210 """Function to add line to existing BokehFigure. 211 212 Args: 213 x_data: list containing x-axis values for line 214 y_data: list containing y_axis values for line 215 legend: string containing line title 216 hover_text: text to display when hovering over lines 217 color: string describing line color 218 marker: string specifying marker, e.g., cross 219 y_axis: identifier for y-axis to plot line against 220 """ 221 if y_axis not in ['default', 'secondary']: 222 raise ValueError('y_axis must be default or secondary') 223 if color == None: 224 color = self.COLORS[self.fig_property['num_lines'] % 225 len(self.COLORS)] 226 if marker == None: 227 marker = self.MARKERS[self.fig_property['num_lines'] % 228 len(self.MARKERS)] 229 self.figure_data.append({ 230 'x_data': x_data, 231 'y_data': y_data, 232 'legend': legend, 233 'hover_text': hover_text, 234 'color': color, 235 'width': 0, 236 'style': 'solid', 237 'marker': marker, 238 'marker_size': marker_size, 239 'shaded_region': None, 240 'y_axis': y_axis 241 }) 242 self.fig_property['num_lines'] += 1 243 244 def generate_figure(self, output_file=None, save_json=True): 245 """Function to generate and save BokehFigure. 246 247 Args: 248 output_file: string specifying output file path 249 """ 250 self.init_plot() 251 two_axes = False 252 for line_idx, line in enumerate(self.figure_data): 253 data_dict = {'x': line['x_data'], 'y': line['y_data']} 254 for key, value in line['hover_text'].items(): 255 data_dict[key] = value 256 source = bokeh.models.ColumnDataSource(data=data_dict) 257 if line['width'] > 0: 258 self.plot.line(x='x', 259 y='y', 260 legend_label=line['legend'], 261 line_width=line['width'], 262 color=line['color'], 263 line_dash=line['style'], 264 name=line['y_axis'], 265 y_range_name=line['y_axis'], 266 source=source) 267 if line['shaded_region']: 268 band_x = line['shaded_region']['x_vector'] 269 band_x.extend(line['shaded_region']['x_vector'][::-1]) 270 band_y = line['shaded_region']['lower_limit'] 271 band_y.extend(line['shaded_region']['upper_limit'][::-1]) 272 self.plot.patch(band_x, 273 band_y, 274 color='#7570B3', 275 line_alpha=0.1, 276 fill_alpha=0.1) 277 if line['error_bars']: 278 data_dict = {'base': numpy.add(line['x_data'], (pow(-1,line_idx)*0.01*line_idx)/self.fig_property['num_lines']), 279 'lower': numpy.subtract(line['y_data'],line['error_bars']), 280 'upper': numpy.add(line['y_data'],line['error_bars']),} 281 source_error = bokeh.models.ColumnDataSource(data=data_dict) 282 self.plot.add_layout( 283 bokeh.models.Whisker(source=source_error, base="base", upper="upper", lower="lower", line_color=line['color']) 284 ) 285 if line['marker'] in self.MARKERS: 286 marker_func = getattr(self.plot, line['marker']) 287 marker_func(x='x', 288 y='y', 289 size=line['marker_size'], 290 legend_label=line['legend'], 291 line_color=line['color'], 292 fill_color=line['color'], 293 name=line['y_axis'], 294 y_range_name=line['y_axis'], 295 source=source) 296 if line['y_axis'] == 'secondary': 297 two_axes = True 298 299 #x-axis formatting 300 self.plot.xaxis.axis_label = self.fig_property['x_label'] 301 self.plot.x_range.range_padding = 0 302 self.plot.xaxis[0].axis_label_text_font_size = self.fig_property[ 303 'axis_label_size'] 304 self.plot.xaxis.major_label_text_font_size = self.fig_property[ 305 'axis_tick_label_size'] 306 #y-axis formatting 307 self.plot.yaxis[0].axis_label = self.fig_property['primary_y_label'] 308 self.plot.yaxis[0].axis_label_text_font_size = self.fig_property[ 309 'axis_label_size'] 310 self.plot.yaxis.major_label_text_font_size = self.fig_property[ 311 'axis_tick_label_size'] 312 default_y_range_renderers = [renderer for renderer in self.plot.renderers if renderer.y_range_name == 'default'] 313 self.plot.y_range = bokeh.models.DataRange1d(name='default', 314 renderers = default_y_range_renderers) 315 if two_axes and 'secondary' not in self.plot.extra_y_ranges: 316 secondary_y_range_renderers = [renderer for renderer in self.plot.renderers if 317 renderer.y_range_name == 'secondary'] 318 319 self.plot.extra_y_ranges = { 320 'secondary': bokeh.models.DataRange1d(name='secondary', 321 renderers=secondary_y_range_renderers) 322 } 323 self.plot.add_layout( 324 bokeh.models.LinearAxis( 325 y_range_name='secondary', 326 axis_label=self.fig_property['secondary_y_label'], 327 axis_label_text_font_size=self.fig_property['axis_label_size']), 'right') 328 # plot formatting 329 self.plot.legend.location = self.fig_property['legend_location'] 330 self.plot.legend.click_policy = 'hide' 331 self.plot.title.text_font_size = self.fig_property['title_size'] 332 self.plot.legend.label_text_font_size = self.fig_property[ 333 'legend_label_size'] 334 335 if output_file is not None: 336 self.save_figure(output_file, save_json) 337 return self.plot 338 339 def load_from_json(self, file_path): 340 with open(file_path, 'r') as json_file: 341 fig_dict = json.load(json_file) 342 self.fig_property = fig_dict['fig_property'] 343 self.figure_data = fig_dict['figure_data'] 344 345 def _save_figure_json(self, output_file): 346 """Function to save a json format of a figure""" 347 figure_dict = collections.OrderedDict(fig_property=self.fig_property, 348 figure_data=self.figure_data) 349 output_file = output_file.replace('.html', '_plot_data.json') 350 with open(output_file, 'w') as outfile: 351 json.dump(wputils.serialize_dict(figure_dict), outfile, indent=4) 352 353 def save_figure(self, output_file, save_json=True): 354 """Function to save BokehFigure. 355 356 Args: 357 output_file: string specifying output file path 358 save_json: flag controlling json outputs 359 """ 360 if save_json: 361 self._save_figure_json(output_file) 362 bokeh.io.output_file(output_file) 363 bokeh.io.save(self.plot) 364 365 @staticmethod 366 def save_figures(figure_array, output_file_path, save_json=True): 367 """Function to save list of BokehFigures in one file. 368 369 Args: 370 figure_array: list of BokehFigure object to be plotted 371 output_file: string specifying output file path 372 """ 373 for idx, figure in enumerate(figure_array): 374 figure.generate_figure() 375 if save_json: 376 json_file_path = output_file_path.replace( 377 '.html', '{}-plot_data.json'.format(idx)) 378 figure._save_figure_json(json_file_path) 379 plot_array = [figure.plot for figure in figure_array] 380 all_plots = bokeh.layouts.column(children=plot_array, 381 sizing_mode='scale_width') 382 bokeh.plotting.output_file(output_file_path) 383 bokeh.plotting.save(all_plots) 384