import beamtime.auxillary as aux import matplotlib.pyplot as plt from matplotlib.ticker import (MultipleLocator, FormatStrFormatter,AutoMinorLocator) from mpl_toolkits.axes_grid.inset_locator import (inset_axes, InsetPosition, mark_inset) import importlib import matplotlib.patches as mpatches from matplotlib.lines import Line2D import matplotlib.lines as mlines from cycler import cycler import itertools def prepare_plot(options={}): ''' Prepares plot based on contents of options['rc_params'] and options['format_params']. rc_params is a dictionary with keyval-pairs corresponding to rcParams in matplotlib format_params will determine the size and aspect ratios of ''' rc_params = options['rc_params'] format_params = options['format_params'] required_options = ['single_column_width', 'double_column_width', 'column_type', 'width_ratio', 'aspect_ratio', 'compress_width', 'compress_height', 'upscaling_factor', 'dpi'] default_options = { 'single_column_width': 8.3, 'double_column_width': 17.1, 'column_type': 'single', 'width_ratio': '1:1', 'aspect_ratio': '1:1', 'compress_width': 1, 'compress_height': 1, 'upscaling_factor': 1.0, 'dpi': 600, } options = aux.update_options(format_params, required_options, default_options) # Reset run commands plt.rcdefaults() # Update run commands if any is passed (will pass an empty dictionary if not passed) update_rc_params(rc_params) width = determine_width(options) height = determine_height(options, width) width, height = scale_figure(options=options, width=width, height=height) fig, ax = plt.subplots(figsize=(width, height), dpi=options['dpi']) return fig, ax def prettify__plot(fig, ax, plot_data, options): required_options = ['plot_kind', 'hide_x_labels', 'hide_y_labels', 'rotation_x_ticks', 'rotation_y_ticks', 'xlabel', 'ylabel', 'yunit', 'xlim', 'ylim', 'x_tick_locators', 'y_tick_locators', 'xticks', 'hide_x_ticks', 'hide_y_ticks', 'hide_x_ticklabels', 'hide_y_ticklabels', 'colours', 'palettes', 'title', 'legend', 'legend_position', 'subplots_adjust', 'text', 'legend_ncol'] default_options = { 'plot_kind': None, # defaults to None, but should be utilised when 'hide_x_labels': False, # Whether x labels should be hidden 'hide_x_ticklabels': False, 'hide_x_ticks': False, 'rotation_x_ticks': 0, 'hide_y_labels': False, # whether y labels should be hidden 'hide_y_ticklabels': False, 'hide_y_ticks': False, 'rotation_y_ticks': 0, 'xlabel': r'$x$ in Na$_{5-x}$FeO$_{4-\delta}$', 'ylabel': 'Formation energy', 'yunit': r'eV', # The unit of the y-values in the curve and bar plots 'xlim': None, 'ylim': None, 'x_tick_locators': [.5, .25], # Major and minor tick locators 'y_tick_locators': [.5, .25], 'xticks': None, 'labels': None, 'colours': None, 'palettes': [('qualitative', 'Dark2_8'), ('qualitative', 'Paired_12')], 'title': None, 'legend': False, 'legend_position': ['lower center', (0.5, -0.1)], # the position of the legend passed as arguments to loc and bbox_to_anchor respectively 'legend_ncol': 1, 'subplots_adjust': [0.1, 0.1, 0.9, 0.9], 'text': None } options = update_options(options=options, required_options=required_options, default_options=default_options) # Set labels on x- and y-axes if not options['hide_y_labels']: ax.set_ylabel(f'{options["ylabel"]} [{options["yunit"]}]') else: ax.set_ylabel('') if not options['hide_x_labels']: ax.set_xlabel(f'{options["xlabel"]}') else: ax.set_xlabel('') # Set multiple locators ax.yaxis.set_major_locator(MultipleLocator(options['y_tick_locators'][0])) ax.yaxis.set_minor_locator(MultipleLocator(options['y_tick_locators'][1])) ax.xaxis.set_major_locator(MultipleLocator(options['x_tick_locators'][0])) ax.xaxis.set_minor_locator(MultipleLocator(options['x_tick_locators'][1])) if options['xticks']: ax.set_xticks(np.arange(plot_data['start'], plot_data['end']+1)) ax.set_xticklabels(options['xticks']) else: ax.set_xticks(np.arange(plot_data['start'], plot_data['end']+1)) ax.set_xticklabels([x/2 for x in np.arange(plot_data['start'], plot_data['end']+1)]) # Hide x- and y- ticklabels if options['hide_y_ticklabels']: ax.tick_params(axis='y', direction='in', which='both', labelleft=False, labelright=False) else: plt.xticks(rotation=options['rotation_x_ticks']) #ax.set_xticklabels(ax.get_xticks(), rotation = options['rotation_x_ticks']) if options['hide_x_ticklabels']: ax.tick_params(axis='x', direction='in', which='both', labelbottom=False, labeltop=False) else: pass #ax.set_yticklabels(ax.get_yticks(), rotation = options['rotation_y_ticks']) # Hide x- and y-ticks: if options['hide_y_ticks']: ax.tick_params(axis='y', direction='in', which='both', left=False, right=False) if options['hide_x_ticks']: ax.tick_params(axis='x', direction='in', which='both', bottom=False, top=False) # Set title if options['title']: ax.set_title(options['title'], fontsize=plt.rcParams['font.size']) # Create legend if ax.get_legend(): ax.get_legend().remove() if options['legend']: # Make palette and linestyles from original parameters if not options['colours']: colours = generate_colours(palettes=options['palettes']) else: colours = itertools.cycle(options['colours']) markers = itertools.cycle(options['markers']) # Create legend active_markers = [] active_labels = [] for label in options['labels']: # Discard next linestyle and colour if label is _ if label == '_': _ = next(colours) _ = next(markers) else: active_markers.append(mlines.Line2D([], [], markeredgecolor=next(colours), color=(1, 1, 1, 0), marker=next(markers))) active_labels.append(label) ax.legend(active_markers, active_labels, frameon=False, loc=options['legend_position'][0], bbox_to_anchor=options['legend_position'][1], ncol=options['legend_ncol']) #fig.legend(handles=patches, loc=options['legend_position'][0], bbox_to_anchor=options['legend_position'][1], frameon=False) # Adjust where the axes start within the figure. Default value is 10% in from the left and bottom edges. Used to make room for the plot within the figure size (to avoid using bbox_inches='tight' in the savefig-command, as this screws with plot dimensions) plt.subplots_adjust(left=options['subplots_adjust'][0], bottom=options['subplots_adjust'][1], right=options['subplots_adjust'][2], top=options['subplots_adjust'][3]) # If limits for x- and y-axes is passed, sets these. if options['xlim'] is not None: ax.set_xlim(options['xlim']) if options['ylim'] is not None: ax.set_ylim(options['ylim']) # Add custom text if options['text']: plt.text(x=options['text'][1][0], y=options['text'][1][1], s=options['text'][0]) return fig, ax def ipywidgets_update(func, plot_data, options={}, **kwargs): for key in kwargs: options[key] = kwargs[key] func(plot_data=plot_data, options=options) def determine_width(options): conversion_cm_inch = 0.3937008 # cm to inch if options['column_type'] == 'single': column_width = options['single_column_width'] elif options['column_type'] == 'double': column_width = options['double_column_width'] column_width *= conversion_cm_inch width_ratio = [float(num) for num in options['width_ratio'].split(':')] width = column_width * width_ratio[0]/width_ratio[1] return width def determine_height(options, width): aspect_ratio = [float(num) for num in options['aspect_ratio'].split(':')] height = width/(aspect_ratio[0] / aspect_ratio[1]) return height def scale_figure(options, width, height): width = width * options['upscaling_factor'] * options['compress_width'] height = height * options['upscaling_factor'] * options['compress_height'] return width, height def update_rc_params(rc_params): ''' Update all passed run commands in matplotlib''' if rc_params: for key in rc_params.keys(): plt.rcParams.update({key: rc_params[key]}) def generate_colours(palettes): # Creates a list of all the colours that is passed in the colour_cycles argument. Then makes cyclic iterables of these. colour_collection = [] for palette in palettes: mod = importlib.import_module("palettable.colorbrewer.%s" % palette[0]) colour = getattr(mod, palette[1]).mpl_colors colour_collection = colour_collection + colour colour_cycle = itertools.cycle(colour_collection) return colour_cycle