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