"""
A simple plotting library used by SSINS.Catalog_Plot. Direct interaction with
these functions is unnecessary.
"""
import numpy as np
from astropy.time import Time
from astropy import units
from astropy.coordinates import Longitude
import warnings
import copy
[docs]def image_plot(fig, ax, data, cmap=None, vmin=None, vmax=None, title='',
xlabel='', ylabel='', midpoint=False, aspect='auto',
cbar_label=None, xticks=None, yticks=None, log=False,
xticklabels=None, yticklabels=None, mask_color='white',
cbar_ticks=None, font_size='medium', symlog=False, linthresh=1,
extent=None, extent_time_format='jd', convert_times=True,
lst_prec=2, extend='neither', alpha=None):
"""
Plots 2-d images. Can do a midpoint normalize and log normalize.
Args:
fig: The fig object to modify
ax: The axis object to modify
data: The data to plot
cmap: Name of a colormap from matplotlib
vmin: The minimum value for the colormap
vmax: The maximum value for the colormap
title: The title for ax
xlabel: The horizontal axis label
ylabel: The vertical axis label
midpoint: Whether to set the midpoint of the colormap to zero (useful for diverging colormaps)
aspect: Adjusts aspect ratio, see ax.imshow documentation in Matplotlib
cbar_label: The label for the colorbar
xticks: The ticks for the horizontal axis
yticks: The ticks for the vertical axis
log: If True, set a logarithmic scale for the colormap
symlog: If True, set a symmetric logarithmic scale for the colormap
linthresh: region for symlog (if enabled) to act linearly to avoid divergence at zero
xticklabels: The labels for the xticks
yticklabels: The labels for the yticks
mask_color: The color for masked data values, if any
cbar_ticks: The tickmarks for the colorbar
font_size: Font size is set globally with this parameter.
extent: Passes to imshow to determine ticks.
extent_time_format: a string specifying the format of the times passed in
the extent keyword.
convert_times: Will convert a JD to UTC or LST in radians to an LST in
hourangle, both using astropy.
lst_prec: Number of sig figs to keep in LST hourangle ticklabel
extend: Whether to extend the colorbar.
alpha: Set a transparency factor
Note for arguments midpoint, log, symlog, linthresh:
* Only one of these arguments can be expressed in the plot (can't have a plot with multiple different colorbar metrics).
* If multiple of these arguments are passed, default is linear (midpoint) followed by log, symlog.
* Linthresh only applies when using symmetrical log (symlog) and will be ignored otherwise.
"""
from matplotlib import colors, cm
class MidpointNormalize(colors.Normalize):
def __init__(self, vmin=None, vmax=None, vcenter=None, clip=False):
self.vcenter = vcenter
super().__init__(vmin, vmax, clip)
def __call__(self, value, clip=None):
# Note also that we must extrapolate beyond vmin/vmax
result, is_scalar = self.process_value(value)
x, y = [self.vmin, self.vcenter, self.vmax], [0, 0.5, 1.]
return np.ma.masked_array(np.interp(value, x, y,
left=-np.inf, right=np.inf),
mask=result.mask)
def inverse(self, value):
y, x = [self.vmin, self.vcenter, self.vmax], [0, 0.5, 1]
return np.interp(value, x, y, left=-np.inf, right=np.inf)
if cmap is None:
cmap = 'plasma'
colormap = copy.copy(getattr(cm, cmap)) # Copy so as not to mutate the global cmap instance
colormap.set_bad(color=mask_color)
# Make sure it does the yticks correctly
if extent is not None:
if (extent_time_format.lower() == 'lst') and (extent[-2] < extent[-1]):
warnings.warn("LSTs appear to cross 24 hrs. Unwrapping. If this is an error, check extent keyword in plot_lib call.")
extent[-1] = extent[-1] - 2 * np.pi
# colorization methods: linear, normalized log, symmetrical log
if midpoint:
cax = ax.imshow(data, cmap=colormap, aspect=aspect, interpolation='none',
norm=MidpointNormalize(vcenter=0, vmin=vmin, vmax=vmax),
extent=extent, alpha=alpha)
elif log:
cax = ax.imshow(data, cmap=colormap, norm=colors.LogNorm(), aspect=aspect,
vmin=vmin, vmax=vmax, interpolation='none',
extent=extent, alpha=alpha)
elif symlog:
cax = ax.imshow(data, cmap=colormap, norm=colors.SymLogNorm(linthresh), aspect=aspect,
vmin=vmin, vmax=vmax, interpolation='none',
extent=extent, alpha=alpha)
else:
cax = ax.imshow(data, cmap=colormap, vmin=vmin, vmax=vmax, aspect=aspect,
interpolation='none', extent=extent, alpha=alpha)
cbar = fig.colorbar(cax, ax=ax, ticks=cbar_ticks, extend=extend)
cbar.set_label(cbar_label, fontsize=font_size)
ax.set_title(title, fontsize=font_size)
ax.set_xlabel(xlabel, fontsize=font_size)
ax.set_ylabel(ylabel, fontsize=font_size)
if extent is None:
set_ticks_labels(ax, xticks, yticks, xticklabels, yticklabels)
elif (xticks is not None) or (yticks is not None):
warnings.warn("Plotting keyword 'extent' has been set alongside xticks "
"or yticks keyword. Using manual settings.")
set_ticks_labels(ax, xticks, yticks, xticklabels, yticklabels)
elif convert_times:
# This case is for when extent is set, manual settings have not been made, and conversion is desired.
# Otherwise just use what came from extent
yticks = ax.get_yticks()
if extent_time_format.lower() == 'jd':
yticklabels = [Time(ytick, format='jd').iso[:-4] for ytick in yticks]
elif extent_time_format.lower() == 'lst':
set_yticklabels = [Longitude(ytick * units.radian).hourangle for ytick in yticks]
set_ticks_labels(ax, xticks, yticks, xticklabels, yticklabels)
cbar.ax.tick_params(labelsize=font_size)
ax.tick_params(labelsize=font_size)
def set_ticks_labels(ax, xticks, yticks, xticklabels, yticklabels):
from matplotlib import ticker
if xticks is not None:
ax.xaxis.set_major_locator(ticker.FixedLocator(xticks))
if yticks is not None:
ax.yaxis.set_major_locator(ticker.FixedLocator(yticks))
if xticklabels is not None:
ax.set_xticklabels(xticklabels)
if yticklabels is not None:
ax.set_yticklabels(yticklabels)
[docs]def hist_plot(fig, ax, data, bins='auto', yscale='log', xscale='linear',
label='', legend=True, title='', xlabel='', ylim=None,
density=False, model_func=None, error_sig=0, model_label='',
color='blue', model_color='orange', alpha=0.5,
font_size='medium', **model_kwargs):
"""
A function that calculates and plots histograms. Can also make models using
the model_func argument, including error shading.
Args:
fig: The figure object to modify
ax: The axis object to modify
data: The data to be histogrammed
bins: The bins for the histogram. See numpy.histogram documentation
yscale ('log' or 'linear'): The scale for the vertical axis. Default is 'log'
xscale ('log' or 'linear'): The scale for the horizontal axis. Default is 'linear'
label: The legend label for the data
legend: If True, show the legend
title: The title for the plot
xlabel: The label for the horizontal axis
ylim: The limits for the vertical axis
density: If True, report probability density instead of counts
model_func: Default is None. Set to a callable function of the bin edges to make a model histogram.
error_sig: Default is 0. If greater than zero, draw error shades out to error_sig sigma (significance).
model_label: The legend label for the model
color: The color of the drawn histogram
model_color: The color of the drawn model (and errors)
alpha: The transparency parameter for the error shades.
model_kwargs: Additional kwargs not included in the above list will be passed to model_func
font_size: Fontsize is set globally here
"""
counts, bins = np.histogram(data, bins=bins, density=density)
counts = np.append(counts, 0)
ax.plot(bins, counts, label=label, drawstyle='steps-post', color=color)
ax.set_xscale(xscale)
ax.set_yscale(yscale)
ax.set_title(title, fontsize=font_size)
ax.set_xlabel(xlabel, fontsize=font_size)
if density:
ax.set_ylabel('Probability Density', fontsize=font_size)
else:
ax.set_ylabel('Counts', fontsize=font_size)
if model_func is not None:
model_prob = model_func(bins, **model_kwargs)
if density:
model_y = model_prob / np.diff(bins)
else:
model_y = model_prob * np.sum(counts)
model_y = np.append(model_y, 0)
if error_sig:
N = np.prod(data.shape)
yerr = np.sqrt(N * model_prob * (1 - model_prob))
if density:
yerr /= (N * np.diff(bins))
yerr = error_sig * np.append(yerr, 0)
ax.fill_between(bins, model_y - yerr, model_y + yerr, alpha=alpha,
step='post', color=model_color)
model_label += ' and %s' % (error_sig) + r'$\sigma$ Uncertainty'
ax.plot(bins, model_y, label=model_label, drawstyle='steps-post', color=model_color)
if legend:
ax.legend(fontsize=font_size)
if ylim is None:
ax.set_ylim([0.9 * np.amin(counts[counts > 0]), 10 * np.amax(counts)])
else:
ax.set_ylim(ylim)
ax.tick_params(labelsize=font_size)