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)
