Source code for EELSFitter.plotting.heatmaps

import math
import warnings
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

from matplotlib import cm
from ..core.spectral_image import round_scientific, trunc


[docs] def plot_heatmap(image, data, dpi=200, title=None, xlabel=r"$\rm{[nm]\;}$", ylabel=r"$\rm{[nm]\;}$", cmap='coolwarm', discrete_colormap=False, sig_cbar=3, color_bin_size=None, sig_ticks=2, npix_xtick=10, npix_ytick=10, scale_ticks=1, tick_int=False, save_as=False, **kwargs): r""" Plots a heatmap for given data input. Parameters ---------- image : SpectralImage :py:meth:`spectral_image.SpectralImage <EELSFitter.core.spectral_image.SpectralImage>` object data : numpy.ndarray, shape=(M,N) Input data for heatmap, but be 2D. dpi : int, optional Set the dpi of the heatmap. The default is 100 title : str, optional Set the title of the heatmap. The default is None. xlabel : str, optional Set the label of the x-axis. Nanometer ([nm]) is assumed as standard scale. The default is '[nm]'. ylabel : str, optional Set the label of the y-axis. Nanometer ([nm]) is assumed as standard scale. The default is '[nm]'. cmap : str, optional Set the colormap of the heatmap. The default is 'coolwarm'. discrete_colormap : bool, optional Enables the heatmap values to be discretised. Best used in conjuction with color_bin_size. The default is False. sig_cbar : int, optional Set the amount of significant numbers displayed in the colorbar. The default is 3. color_bin_size : float, optional Set the size of the bins used for discretisation. Best used in conjuction discrete_colormap. The default is None. sig_ticks : int, optional Set the amount of significant numbers displayed in the ticks. The default is 2. npix_xtick : float, optional Display a tick per n pixels in the x-axis. Note that this value can be a float. The default is 10. npix_ytick : float, optional Display a tick per n pixels in the y-axis. Note that this value can be a float. The default is 10. scale_ticks : float, optional Change the scaling of the numbers displayed in the ticks. Nanometer ([nm]) is assumed as standard scale adjust scaling from there. The default is 1. tick_int : bool, optional Set whether you only want the ticks to display as integers instead of floats. The default is False. save_as : str, optional Set the location and name for the heatmap to be saved to. The default is False. **kwargs : dictionary Additional keyword arguments. Returns ------- fig: Seaborn.figure.Figure """ fig = plt.figure(dpi=dpi) if title is None: if image.name is not None: plt.title(image.name) else: plt.title(title) if 'mask' in kwargs: mask = kwargs['mask'] if mask.all(): warnings.warn("Mask all True: no values to plot.") return else: mask = np.zeros(data.shape).astype('bool') # Create the discretization using the given discretized data if discrete_colormap: unique_data_points = np.unique(data[~mask]) if 'vmax' in kwargs: if len(unique_data_points[unique_data_points > kwargs['vmax']]) > 0: unique_data_points = unique_data_points[unique_data_points <= kwargs['vmax']] if 'vmin' in kwargs: if len(unique_data_points[unique_data_points < kwargs['vmin']]) > 0: unique_data_points = unique_data_points[unique_data_points >= kwargs['vmin']] if color_bin_size is None: if len(unique_data_points) == 1: color_bin_size = 1 else: color_bin_size = np.nanpercentile(unique_data_points[1:] - unique_data_points[:-1], 30) n_colors = round((np.nanmax(unique_data_points) - np.nanmin(unique_data_points)) / color_bin_size + 1) cmap = cm.get_cmap(cmap, n_colors) spacing = color_bin_size / 2 kwargs['vmax'] = np.max(unique_data_points) + spacing kwargs['vmin'] = np.min(unique_data_points) - spacing # Creat the heatmap if image.pixel_size is not None: ax = sns.heatmap(data, cmap=cmap, square=True, **kwargs) xticks, yticks, xticks_labels, yticks_labels = get_ticks(image, sig_ticks, npix_xtick, npix_ytick, scale_ticks, tick_int) ax.xaxis.set_ticks(xticks) ax.yaxis.set_ticks(yticks) ax.set_xticklabels(xticks_labels, rotation=0) ax.set_yticklabels(yticks_labels) else: ax = sns.heatmap(data, **kwargs) ax.set_xlabel(xlabel) ax.set_ylabel(ylabel) # Create the discretized colorbar from the discretized data colorbar = ax.collections[0].colorbar if discrete_colormap: if data.dtype == int: colorbar.set_ticks(unique_data_points) else: colorbar.set_ticks(unique_data_points) cbar_ticks_labels = [] for tick in unique_data_points: if tick >= 1: cbar_ticks_labels.append(round_scientific(tick, sig_cbar + len(str(abs(int(math.floor(tick))))))) else: cbar_ticks_labels.append(round_scientific(tick, sig_cbar)) colorbar.ax.set_yticklabels(cbar_ticks_labels) # Adds equal or greater than symbol to max color value if 'vmax' in kwargs: if np.nanmax(data[~mask]) > kwargs['vmax']: cbar_ticks = colorbar.ax.get_yticklabels() loc = -1 if discrete_colormap: loc = np.max(np.argwhere(colorbar.ax.get_yticks() <= kwargs['vmax'] - spacing)) cbar_ticks[loc] = r'$\geq$' + cbar_ticks[loc].get_text() colorbar.ax.set_yticklabels(cbar_ticks) # Adds equal or less than symbol to min color value if 'vmin' in kwargs: if np.nanmin(data[~mask]) < kwargs['vmin']: cbar_ticks = colorbar.ax.get_yticklabels() loc = 0 if discrete_colormap: loc = np.min(np.argwhere(colorbar.ax.get_yticks() >= kwargs['vmin'] + spacing)) cbar_ticks[loc] = r'$\leq$' + cbar_ticks[loc].get_text() colorbar.ax.set_yticklabels(cbar_ticks) if save_as: if type(save_as) != str: if image.name is not None: save_as = image.name if 'mask' in kwargs: save_as += '_masked' save_as += '.pdf' plt.savefig(save_as, bbox_inches='tight') return fig
[docs] def get_ticks(image, sig_ticks=2, npix_xtick=10, npix_ytick=10, scale_ticks=1, tick_int=False): r""" Sets the proper tick labels and tick positions for the heatmap plots. Parameters ---------- sig_ticks : int, optional Set the amount of significant numbers displayed in the ticks. The default is 2. npix_xtick : float, optional Display a tick per n pixels in the x-axis. Note that this value can be a float. The default is 10. npix_ytick : float, optional Display a tick per n pixels in the y-axis. Note that this value can be a float. The default is 10. scale_ticks : float, optional Change the scaling of the numbers displayed in the ticks. Microns ([\u03BCm]) are assumed as standard scale, adjust scaling from there. The default is 1. tick_int : bool, optional Set whether you only want the ticks to display as integers instead of floats. The default is False. Returns ------- xticks : numpy.ndarray, shape=(M,) Array of the xticks positions. yticks : numpy.ndarray, shape=(M,) Array of the yticks positions. xticks_labels : numpy.ndarray, shape=(M,) Array with strings of the xtick labels. yticks_labels : numpy.ndarray, shape=(M,) Array with strings of the ytick labels. """ image.calc_axes() xticks = np.arange(0, image.x_axis.shape[0], npix_xtick) yticks = np.arange(0, image.y_axis.shape[0], npix_ytick) if tick_int == True: xticks_labels = (xticks * round_scientific(image.pixel_size[1] * scale_ticks, sig_ticks)).astype(int) yticks_labels = (yticks * round_scientific(image.pixel_size[0] * scale_ticks, sig_ticks)).astype(int) else: xticks_labels = trunc(xticks * round_scientific(image.pixel_size[1] * scale_ticks, sig_ticks), sig_ticks) yticks_labels = trunc(yticks * round_scientific(image.pixel_size[0] * scale_ticks, sig_ticks), sig_ticks) return xticks, yticks, xticks_labels, yticks_labels