# -*- coding: utf-8 -*- """ Utilities for generating static image and line plots of near-publishable quality Created on Thu May 05 13:29:12 2016 @author: Suhas Somnath, Chris R. Smith """ from __future__ import division, print_function, absolute_import, \ unicode_literals import inspect import sys from numbers import Number import matplotlib as mpl import numpy as np from dask import array as da from matplotlib import pyplot as plt from mpl_toolkits.axes_grid1 import ImageGrid from sidpy.base.num_utils import get_exponent from sidpy.viz.plot_utils.misc import get_plot_grid_size, set_tick_font_size from sidpy.viz.plot_utils.cmap import default_cmap if sys.version_info.major == 3: unicode = str def plot_map(axis, img, show_xy_ticks=True, show_cbar=True, x_vec=None, y_vec=None, num_ticks=4, stdevs=None, cbar_label=None, tick_font_size=None, infer_aspect=False, **kwargs): """ Plots an image within the given axis with a color bar + label and appropriate X, Y tick labels. This is particularly useful to get readily interpretable plots for papers Parameters ---------- axis : matplotlib.axes.Axes object Axis to plot this image onto img : 2D numpy array with real values Data for the image plot show_xy_ticks : bool, Optional, default = None, shown unedited Whether or not to show X, Y ticks show_cbar : bool, optional, default = True Whether or not to show the colorbar x_vec : 1-D array-like or Number, optional if an array-like is provided, these will be used for the tick values on the X axis if a Number is provided, this will serve as an extent for tick values in the X axis. For example x_vec=1.5 would cause the x tick labels to range from 0 to 1.5 y_vec : 1-D array-like or Number, optional if an array-like is provided - these will be used for the tick values on the Y axis if a Number is provided, this will serve as an extent for tick values in the Y axis. For example y_vec=225 would cause the y tick labels to range from 0 to 225 num_ticks : unsigned int, optional, default = 4 Number of tick marks on the X and Y axes stdevs : unsigned int (Optional. Default = None) Number of standard deviations to consider for plotting. If None, full range is plotted. cbar_label : str, optional, default = None Labels for the colorbar. Use this for something like quantity (units) tick_font_size : unsigned int, optional, default = None Font size to apply to x, y, colorbar ticks and colorbar label infer_aspect : bool, Optional. Default = False Whether or not to adjust the aspect ratio of the image based on the provided x_vec and y_vec The values of x_vec and y_vec will be assumed to have the same units. kwargs : dictionary Anything else that will be passed on to matplotlib.pyplot.imshow Returns ------- im_handle : handle to image plot handle to image plot cbar : handle to color bar handle to color bar Note ---- The origin of the image will be set to the lower left corner. Use the kwarg 'origin' to change this """ if not isinstance(axis, mpl.axes.Axes): raise TypeError('axis must be a matplotlib.axes.Axes object') if not isinstance(img, (np.ndarray, da.core.Array)): raise TypeError('img should be a numpy array') if not img.ndim == 2: raise ValueError('img should be a 2D array') if not isinstance(show_xy_ticks, bool): raise TypeError('show_xy_ticks should be a boolean value') if not isinstance(show_cbar, bool): raise TypeError('show_cbar should be a boolean value') # checks for x_vec and y_vec are done below if num_ticks is not None: if not isinstance(num_ticks, int): raise TypeError('num_ticks should be a whole number') if num_ticks < 2: raise ValueError('num_ticks should be at least 2') if tick_font_size is not None: if not isinstance(tick_font_size, Number): raise TypeError('tick_font_size must be a whole number') if tick_font_size < 0: raise ValueError('tick_font_size must be a whole number') if stdevs is not None: if not isinstance(stdevs, Number): raise TypeError('stdevs should be a Number') data_mean = np.mean(img) data_std = np.std(img) kwargs.update({'clim': [data_mean - stdevs * data_std, data_mean + stdevs * data_std]}) kwargs.update({'origin': kwargs.pop('origin', 'lower')}) if show_cbar: if np.isnan(img).any(): _img = img[np.where(~np.isnan(img))] y_exp = get_exponent(np.squeeze(_img)) else: y_exp = get_exponent(np.squeeze(img)) z_suffix = '' if y_exp < -2 or y_exp > 3: img = np.squeeze(img) / 10 ** y_exp z_suffix = ' x $10^{' + str(y_exp) + '}$' assert isinstance(show_xy_ticks, bool) ######################################################################################################## def set_ticks_for_axis(tick_vals, is_x): if is_x: tick_vals_var_name = 'x_vec' tick_set_func = axis.set_xticks tick_labs_set_func = axis.set_xticklabels else: tick_vals_var_name = 'y_vec' tick_set_func = axis.set_yticks tick_labs_set_func = axis.set_yticklabels img_axis = int(is_x) img_size = img.shape[img_axis] chosen_ticks = np.linspace(0, img_size - 1, num_ticks, dtype=int) if tick_vals is not None: if isinstance(tick_vals, (int, float)): if tick_vals > 0.01: tick_labs = [str(np.round(ind * tick_vals / img_size, 2)) for ind in chosen_ticks] else: tick_labs = ['{0:.1E}'.format(ind * tick_vals / img_size) for ind in chosen_ticks] print(tick_labs) tick_vals = np.linspace(0, tick_vals, img_size) else: if not isinstance(tick_vals, (np.ndarray, list, tuple, range, da.core.Array)) or \ len(tick_vals) != img_size: raise ValueError( '{} should be array-like with shape equal to axis {} of img'.format(tick_vals_var_name, img_axis)) if np.max(tick_vals) > 0.01: tick_labs = [str(np.round(tick_vals[ind], 2)) for ind in chosen_ticks] else: tick_labs = ['{0:.1E}'.format(tick_vals[ind]) for ind in chosen_ticks] else: tick_labs = [str(ind) for ind in chosen_ticks] tick_set_func(chosen_ticks) tick_labs_set_func(tick_labs) if tick_font_size is not None: set_tick_font_size(axis, tick_font_size) return tick_vals ######################################################################################################## if show_xy_ticks is True or x_vec is not None: x_vec = set_ticks_for_axis(x_vec, True) else: axis.set_xticks([]) if show_xy_ticks is True or y_vec is not None: y_vec = set_ticks_for_axis(y_vec, False) else: axis.set_yticks([]) if infer_aspect: # Aspect ratio determined by this function will take precedence. _ = kwargs.pop('infer_aspect', None) """ At this stage, if x_vec and y_vec are not None, they should be arrays. This will be very useful when one dimension is coarsely sampled while another is finely sampled and we want to visualize the image with the physically correct aspect ratio. This CANNOT be performed automatically due to potentially incompatible units which are unknown to this func. """ if x_vec is not None or y_vec is not None: x_range = x_vec.max() - x_vec.min() y_range = y_vec.max() - y_vec.min() kwargs.update({'aspect': (y_range / x_range) * (img.shape[1] / img.shape[0])}) im_handle = axis.imshow(img, **kwargs) cbar = None if not isinstance(show_cbar, bool): show_cbar = False if show_cbar: cbar = plt.colorbar(im_handle, ax=axis, orientation='vertical', fraction=0.046, pad=0.04, use_gridspec=True) # cbar = axis.cbar_axes[count].colorbar(im_handle) if cbar_label is not None: if not isinstance(cbar_label, (str, unicode)): raise TypeError('cbar_label should be a string') if tick_font_size is not None: cbar.set_label(cbar_label + z_suffix) else: cbar.set_label(cbar_label + z_suffix, fontsize=tick_font_size) else: if z_suffix != '': cbar.set_label(z_suffix) if tick_font_size is not None: cbar.ax.tick_params(labelsize=tick_font_size) return im_handle, cbar def plot_map_stack(map_stack, num_comps=9, stdevs=2, color_bar_mode=None, evenly_spaced=False, reverse_dims=False, subtitle='Component', title=None, colorbar_label='', fig_mult=(5, 5), pad_mult=(0.1, 0.07), x_label=None, y_label=None, title_yoffset=None, title_size=None, **kwargs): """ Plots the provided stack of maps Parameters ------------- map_stack : 3D real numpy array structured as [component, rows, cols] num_comps : int, Optional Number of components to plot stdevs : int, Optional Number of standard deviations to consider for plotting. Set to None if no clipping is desired color_bar_mode : String, Optional Options are None, single or each. Default None evenly_spaced : bool, Optional. Default = False If set to True - The slices / component will be selected at intervals from the first to last If set to False - The first ``num_comps`` images will be plotted instead reverse_dims : bool, Optional. Default = False Set this to True to accept data structured as [rows, cols, component] subtitle : String or list of strings The titles for each of the plots. If a single string is provided, the plot titles become ['title 01', title 02', ...]. if a list of strings (equal to the number of components) are provided, these are used instead. title : str, Optinal Title for the plot grid that will appear at the top colorbar_label : str, Optional label for colorbar. Default is an empty string. fig_mult : length 2 array_like of uints Size multipliers for the figure. Figure size is calculated as (num_rows*`fig_mult[0]`, num_cols*`fig_mult[1]`). Default (4, 4) pad_mult : tuple, list, array-like, Optional Array-like of floats of length 2. Multipliers for the axis padding between plots in the stack. Padding is calculated as (pad_mult[0]*fig_mult[1], pad_mult[1]*fig_mult[0]) for the width and height padding respectively. Default (0.1, 0.07) x_label : str, Optional X Label for all plots y_label : (optional) String Y label for all plots title_yoffset : float Offset to move the figure title vertically in the figure title_size : float Size of figure title kwargs : dictionary Keyword arguments to be passed to either matplotlib.pyplot.figure, mpl_toolkits.axes_grid1.ImageGrid, or pyUSID.viz.plot_utils.plot_map. See specific function documentation for the relavent options. Returns --------- fig, axes """ # plt.rcParams["mpl_toolkits.legacy_colorbar"] = False if not isinstance(map_stack, (np.ndarray, da.core.Array)) or not map_stack.ndim == 3: raise TypeError('map_stack should be a 3 dimensional array arranged as [component, row, col]') if num_comps is None: num_comps = 4 # Default else: if not isinstance(num_comps, int) or num_comps < 1: raise TypeError('num_comps should be a positive integer') for var, var_name in zip([title, colorbar_label, color_bar_mode, x_label, y_label], ['title', 'colorbar_label', 'color_bar_mode', 'x_label', 'y_label']): if var is not None: if not isinstance(var, (str, unicode)): raise TypeError(var_name + ' should be a string') if title is None: title = '' if colorbar_label is None: colorbar_label = '' if x_label is None: x_label = '' if y_label is None: y_label = '' if color_bar_mode not in [None, 'single', 'each']: raise ValueError('color_bar_mode must be either None, "single", or "each"') for var, var_name in zip([stdevs, title_yoffset, title_size], ['stdevs', 'title_yoffset', 'title_size']): if var is not None: if not isinstance(var, Number) or var <= 0: raise TypeError(var_name + ' of value: {} should be a number > 0'.format(var)) for var, var_name in zip([evenly_spaced, reverse_dims], ['evenly_spaced', 'reverse_dims']): if not isinstance(var, bool): raise TypeError(var_name + ' should be a bool') for var, var_name in zip([fig_mult, pad_mult], ['fig_mult', 'pad_mult']): if not isinstance(var, (list, tuple, np.ndarray, da.core.Array)) or len(var) != 2: raise TypeError(var_name + ' should be a tuple / list / numpy array of size 2') if not np.all([x > 0 and isinstance(x, Number) for x in var]): raise ValueError(var_name + ' should contain positive numbers') if reverse_dims: map_stack = np.transpose(map_stack, (2, 0, 1)) num_comps = abs(num_comps) num_comps = min(num_comps, map_stack.shape[0]) if evenly_spaced: chosen_pos = np.linspace(0, map_stack.shape[0] - 1, num_comps, dtype=int) else: chosen_pos = np.arange(num_comps, dtype=int) if isinstance(subtitle, list): if len(subtitle) > num_comps: # remove additional subtitles subtitle = subtitle[:num_comps] elif len(subtitle) < num_comps: # add subtitles subtitle += ['Component' + ' ' + str(x) for x in range(len(subtitle), num_comps)] else: if not isinstance(subtitle, str): subtitle = 'Component' subtitle = [subtitle + ' ' + str(x) for x in chosen_pos] fig_h, fig_w = fig_mult p_rows, p_cols = get_plot_grid_size(num_comps) if p_rows * p_cols < num_comps: p_cols += 1 pad_w, pad_h = pad_mult ''' Set defaults for kwargs to the figure creation and extract any non-default values from current kwargs ''' figkwargs = dict() if sys.version_info.major == 3: inspec_func = inspect.getfullargspec else: inspec_func = inspect.getargspec for key in inspec_func(plt.figure).args: if key in kwargs: figkwargs.update({key: kwargs.pop(key)}) fig = plt.figure(figsize=(p_cols * fig_w, p_rows * fig_h), **figkwargs) ''' Set defaults for kwargs to the ImageGrid and extract any non-default values from current kwargs ''' igkwargs = {'cbar_pad': '1%', 'cbar_size': '5%', 'cbar_location': 'right', 'direction': 'row', 'share_all': False, 'aspect': True, 'label_mode': 'L'} # 'add_all': True} for key in igkwargs.keys(): if key in kwargs: igkwargs.update({key: kwargs.pop(key)}) axes = ImageGrid(fig=fig, rect=111, nrows_ncols=(p_rows, p_cols), cbar_mode=color_bar_mode, axes_pad=(pad_w * fig_w, pad_h * fig_h), **igkwargs) try: fig.canvas.set_window_title(title) except: fig.canvas.manager.set_window_title(title) # These parameters have not been easy to fix: if title_yoffset is None: title_yoffset = 0.9 if title_size is None: title_size = 16 + (p_rows + p_cols) fig.suptitle(title, fontsize=title_size, y=title_yoffset) # plt.rcParams["mpl_toolkits.legacy_colorbar"] = False for count, index, curr_subtitle in zip(range(chosen_pos.size), chosen_pos, subtitle): im, im_cbar = plot_map(axes[count], map_stack[index], stdevs=stdevs, show_cbar=False, **kwargs) axes[count].set_title(curr_subtitle) if color_bar_mode == 'each': cb = axes[count].cax.colorbar(im) if count % p_cols == p_cols-1: cb.set_label(colorbar_label) if count % p_cols == 0: axes[count].set_ylabel(y_label) if count >= (p_rows - 1) * p_cols: axes[count].set_xlabel(x_label) # With cbar_mode="single", cax attribute of all axes are identical. if color_bar_mode == 'single': cb = axes[0].cax.colorbar(im) cb.set_label(colorbar_label) return fig, axes