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