def plot(self, **kwargs): style = kwargs.pop('style', 'errorbar') kwargs['label'] = get_lookup_val('label', kwargs.get('label')) # Plot object on this axis axis_name = kwargs.pop('axis', 'ax') try: ax = getattr(self, axis_name) except AttributeError as e: log.critical('The axis name {0} does not exist.'.format(axis_name)) log.critical(e) raise if kwargs['color'] == 'auto': kwargs['color'] = next(self.auto_colors) if kwargs['edgecolor'] == 'auto': kwargs['edgecolor'] = kwargs['color'] kwargs['color'] = get_lookup_val('color', kwargs.get('color')) # Facecolor is always set using the color argument. kwargs['facecolor'] = kwargs['color'] kwargs['edgecolor'] = get_lookup_val('color', kwargs.get('edgecolor')) if style == 'errorbar': artist = plot_errorbar(ax=ax, **kwargs) elif style == 'band': artist = plot_band(ax=ax, **kwargs) elif style == 'histo': artist = plot_histo(ax=ax, **kwargs) elif style == 'line': artist = plot_line(ax=ax, **kwargs) elif style == 'errorlines': artist = plot_errorlines(ax=ax, **kwargs) elif style == 'heatmap': # special case for z scale and lims in heatmaps since they have to be set by the object instead of the axis. kwargs['z_log'] = self.z_log kwargs['z_lims'] = self.z_lims artist = plot_heatmap(ax=self.ax, **kwargs) self.colorbar_mappable = artist else: raise ValueError('Style {0} not supported.'.format(style)) self._ids.append(kwargs['id']) self._legend_handles.append(artist) self._legend_labels.append(kwargs['label']) return artist
def __call__(self, config): print config.get('data_lims', []) if any(id == 'all' for id, kwargs in config.get('data_lims', [])): all_kwargs = [kwargs for id, kwargs in config['data_lims'] if id == 'all'][0] config['data_lims'] = [(id, all_kwargs) for id in config['objects'].keys() if (id != '_default')] for id, kwargs in config['data_lims']: if id not in config['objects']: raise ValueError('Requested id {} not found.'.format(id)) if not 'obj' in config['objects'][id]: continue log.info('Rebinning id {0}'.format(id)) obj = config['objects'][id]['obj'] min_val = float(get_lookup_val('data_lims', kwargs.get('min', None))) max_val = float(get_lookup_val('data_lims', kwargs.get('max', None))) if isinstance(obj, ROOT.TH1): config['objects'][id]['obj'] = rebin_histo(obj, min_val, max_val) elif isinstance(obj, ROOT.TGraph): config['objects'][id]['obj'] = rebin_tgraph(obj, min_val, max_val)
def __init__(self, histos=None, **kwargs): super(Plot, self).__init__(**kwargs) if not kwargs['add_subplot']: self.ax = self.fig.add_subplot(111) self.ax1 = None else: self.ax = plt.subplot2grid((4, 1), (0, 0), rowspan=3) self.ax1 = plt.subplot2grid((4, 1), (3, 0), rowspan=1) self.histos = histos self.x_lims = [get_lookup_val('x_lims', val) for val in kwargs.pop('x_lims', [None, None])] self.x_lims = [any2float(v) for v in self.x_lims] self.y_lims = [get_lookup_val('y_lims', val) for val in kwargs.pop('y_lims', [None, None])] self.y_lims = [any2float(v) for v in self.y_lims] self.y_subplot_lims = kwargs.pop('y_subplot_lims', (None, None)) self.y_subplot_lims = [any2float(v) for v in self.y_subplot_lims] self.z_lims = kwargs.pop('z_lims', (None, None)) self.z_lims = [any2float(v) for v in self.z_lims] self.margin = kwargs.pop('margin', 0.1) self.x_log = kwargs.pop('x_log', False) self.y_log = kwargs.pop('y_log', False) self.z_log = kwargs.pop('z_log', False) self.x_axis_formatter = kwargs.pop('x_axis_formatter', 'scalar') self.y_axis_formatter = kwargs.pop('y_axis_formatter', 'scalar') self.z_axis_formatter = kwargs.pop('z_axis_formatter', 'scalar') self.x_label = get_lookup_val('x_label', kwargs.pop('x_label', '')) self.y_label = get_lookup_val('y_label', kwargs.pop('y_label', '')) self.y_subplot_label = get_lookup_val('y_label', kwargs.pop('y_subplot_label', '')) self.z_label = get_lookup_val('z_label', kwargs.pop('z_label', '')) self.show_legend = kwargs.pop('show_legend', True) self.combine_legend_entries = kwargs.pop('combine_legend_entries', []) self.legend_loc = kwargs.pop('legend_loc', 'best') self.legend_ncol = kwargs.pop('legend_ncol', 1) self.legend_bbox_anchor = kwargs.pop('legend_bbox_anchor', None) self.texts = kwargs.pop('ax_texts', []) self.vlines = kwargs.pop('ax_vlines', []) self.hlines = kwargs.pop('ax_hlines', []) self.auto_colors = itertools.cycle(matplotlib.rcParams['axes.color_cycle']) self.colorbar_mappable = None # helpers for legend labels/artists self._ids = [] self._legend_handles = [] self._legend_labels = []
def finish(self): # Add colorbar if there is a mappable if self.colorbar_mappable: cb = self.fig.colorbar(self.colorbar_mappable, ax=self.ax) cb.ax.minorticks_on() cb.solids.set_rasterized(True) if self.z_label: cb.set_label(self.z_label) # Add axis texts for text in self.texts: text = get_lookup_val('ax_texts', text) default_text_kwargs = {'x': 0.05, 'y': 0.95, 'va': 'top', 'ha': 'left'} text_kwargs = parse_query(text) default_text_kwargs.update(text_kwargs) ax_name = default_text_kwargs.pop('axis', 'ax') s = default_text_kwargs.pop('s') r = re.compile(r'\$([^$]*)\$') s = r.sub(lambda m: m.group().replace('.', '.\!'), s) try: cax = getattr(self, ax_name) except AttributeError as e: log.critical('The axis name {0} does not exist.'.format(ax_name)) log.critical(e) raise cax.text(s=s, transform=cax.transAxes, **default_text_kwargs) # Add horizontal lines to ax for hline_kwargs in self.hlines: ax_name = hline_kwargs.pop('axis', 'ax') try: cax = getattr(self, ax_name) except AttributeError as e: log.critical('The axis name {0} does not exist.'.format(ax_name)) log.critical(e) raise cax.axhline(**hline_kwargs) # Add vertical lines to ax for vline_kwargs in self.vlines: ax_name = vline_kwargs.pop('axis', 'ax') try: cax = getattr(self, ax_name) except AttributeError as e: log.critical('The axis name {0} does not exist.'.format(ax_name)) log.critical(e) raise cax.axvline(**vline_kwargs) # a specified position of the label can be set via label?json_dict x_label_kwargs = {'position': (1.0, 0.0), 'ha': 'right', 'va': 'top'} x_label, user_x_label_kwargs = parse_optionstring(self.x_label) x_label_kwargs.update(user_x_label_kwargs) if self.ax1: self.ax1.set_xlabel(x_label, **x_label_kwargs) else: self.ax.set_xlabel(x_label, **x_label_kwargs) y_label_kwargs = {'position': (0.0, 1.0), 'ha': 'right', 'va': 'bottom'} y_label, user_y_label_kwargs = parse_optionstring(self.y_label) y_label_kwargs.update(user_y_label_kwargs) self.ax.set_ylabel(y_label, **y_label_kwargs) if self.ax1: y_subplot_label, y_subplot_label_kwargs = parse_optionstring(self.y_subplot_label) self.ax1.set_ylabel(y_subplot_label, **y_subplot_label_kwargs) if self.x_log: self.ax.set_xscale('log') if self.x_axis_formatter == 'scalar': xfmt = ScalarFormatter() self.ax.xaxis.set_major_formatter(xfmt) elif self.x_axis_formatter == 'scalar2': xfmt = ScalarFormatter() self.ax.xaxis.set_minor_formatter(plt.FuncFormatter(log_locator_filter)) self.ax.xaxis.set_major_formatter(xfmt) if self.ax1: self.ax1.set_xscale('log') if self.x_axis_formatter == 'scalar': xfmt = ScalarFormatter() self.ax1.xaxis.set_major_formatter(xfmt) elif self.x_axis_formatter == 'scalar2': xfmt = ScalarFormatter() self.ax1.xaxis.set_minor_formatter(plt.FuncFormatter(log_locator_filter)) self.ax1.xaxis.set_major_formatter(xfmt) else: self.ax.set_xscale('linear') if self.y_log: self.ax.set_yscale('log', nonposy='clip') if self.y_axis_formatter == 'scalar': xfmt = ScalarFormatter() self.ax.yaxis.set_major_formatter(xfmt) elif self.y_axis_formatter == 'scalar2': xfmt = ScalarFormatter() self.ax.yaxis.set_major_formatter(xfmt) self.ax.yaxis.set_minor_formatter(plt.FuncFormatter(log_locator_filter)) else: self.ax.set_yscale('linear') # By default set a 10% margin plot_tools.set_margin(margin=self.margin) self.ax.set_ylim(ymin=self.y_lims[0], ymax=self.y_lims[1]) self.ax.set_xlim(xmin=self.x_lims[0], xmax=self.x_lims[1]) if self.show_legend: # handles, labels = self.ax.get_legend_handles_labels() no_legend_ids = ['nolegend', '_nolegend_', '__nolegend__','none', ''] # handles = self._legend_handles # labels = self._legend_labels # TODO combine legend entries for id, id2 in self.combine_legend_entries: log.debug('Combining legend entries {0} and {1}'.format(id, id2)) if id in self._ids and id2 in self._ids: self._legend_handles[self._ids.index(id)] = (self._legend_handles[self._ids.index(id2)],self._legend_handles[self._ids.index(id)]) leg_entry_dict = OrderedDict(zip(self._legend_labels, self._legend_handles)) for key in leg_entry_dict.keys(): if key.lower() in no_legend_ids: del leg_entry_dict[key] if leg_entry_dict: labels, handles = zip(*leg_entry_dict.items()) if 'outside' in self.legend_loc: if self.legend_bbox_anchor: bbox_to_anchor = self.legend_bbox_anchor else: bbox_to_anchor = (1, 1) self.legend_loc = self.legend_loc.replace('outside', '').strip() else: bbox_to_anchor = None legend = self.ax.legend(handles, labels, loc=self.legend_loc, ncol=self.legend_ncol, bbox_to_anchor=bbox_to_anchor) legend.get_frame().set_alpha(0.0) # [obj.set_rasterized(True) for obj in legend.get_patches()] # for obj in legend.get_patches(): # obj.set_rasterized(True) # legend.draw() # self.ax.legend_ = None # self.ax.add_artist(legend) else: log.debug('Omit legend since all labels are empty.') if self.ax1: self.ax1.set_ylim(ymin=self.y_subplot_lims[0], ymax=self.y_subplot_lims[1]) plt.setp(self.ax.get_xticklabels(), visible=False) plt.setp(self.ax.get_xticklabels(minor=True), visible=False) # self.ax1.set_xscale(self.ax.get_xscale()) self.ax1.set_xlim(self.ax.get_xlim()) plt.subplots_adjust(hspace=0.15) callbacks.trigger('after_plot', plt=self) self.save_fig() plt.close(self.fig)