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