"""
Module for the BaseExtendableHeatmap class, which provides the basic
functionality of the extendable heatmap plotter.
"""
from __future__ import division
import numpy as np
from matplotlib import colors
import seaborn as sns
from lib5c.plotters.colormaps import get_colormap
from lib5c.plotters.extendable.extendable_figure import ExtendableFigure
[docs]class BaseExtendableHeatmap(ExtendableFigure):
"""
Minimal implementation of an ExtendableFigure organized around a contact
frequency heatmap.
The heatmap is plotted using ``plt.imshow()`` using data from the ``array``
attribute, colored using the other attributes, and the resulting axis is
accessible at ``h['root']`` where ``h`` is the ExtendableHeatmap instance.
New axes can be added to the margins of the heatmap by calling
``add_margin_ax()``. The axis of the new axis that is parallel to the
heatmap will have its limits set to match the ``grange_x`` or ``grange_y``
attributes. This allows plotting features on the margin axes using genomic
coordinates instead of having to convert to pixel coordinates.
The root heatmap axis is still kept in units of pixels. To make drawing on
this axis easier, this class provides ``transform_feature()``, which will
transform a genomic feature dict into heatmap pixel units.
Attributes
----------
array : np.ndarray
Array of values to plot in the heatmap. Must be square.
grange_x : dict
The genomic range represented by the x-axis of this heatmap. The dict
should have the form::
{
'chrom': str,
'start': int,
'end': int
}
grange_y : dict, optional
The genomic range represented by the y-axis of this heatmap. If None,
the heatmap is assumed to be symmetric.
colorscale : tuple of float
The (min, max) of the color range to plot the values in the array with.
colormap : str or matplotlib colormap or dict of colors
The colormap to use when drawing the heatmap. Strings will be passed to
`lib5c.plotters.colormaps.get_colormap()`. If `array` contains strings,
pass a dict mapping those strings to colors.
norm : matplotlib.colors.Normalize, optional
Pass an instance of ``matplotlib.colors.Normalize`` to apply this
normalization to the heatmap and colorbar.
"""
def __init__(self, array, grange_x, grange_y=None, colorscale=None,
colormap='obs', norm=None):
super(BaseExtendableHeatmap, self).__init__()
self.grange_x = grange_x
self.grange_y = grange_y if grange_y is not None else grange_x
self.categorical = any(c in str(array.dtype) for c in ['U', 'S'])
if self.categorical:
keys = list(sorted(np.unique(array)))
if colormap is None or type(colormap) == str:
keys = list(sorted(set(keys) - {'', 'background'}))
colormap = dict(zip(keys, sns.color_palette("husl", len(keys))))
colormap[''] = '#333333'
colormap['background'] = '#333333'
else:
for key in keys:
if key not in colormap:
colormap[key] = '#333333'
self.cmap = colors.ListedColormap(
[colormap[k] for k in sorted(colormap.keys())])
self.color_keys = list(sorted(colormap.keys()))
for k in self.color_keys:
array = np.copy(array)
array[array == k] = self.color_keys.index(k)
array = array.astype(int)
self.vmin = 0
self.vmax = len(self.color_keys) - 1
else:
if colorscale is None:
colorscale = (np.nanmin(array), np.nanmax(array))
self.vmin = colorscale[0]
self.vmax = colorscale[1]
self.cmap = get_colormap(colormap) if type(colormap) == str \
else colormap
self.color_keys = None
self.array = array
self.im = self['root'].imshow(
self.array, interpolation='none', vmin=self.vmin, vmax=self.vmax,
cmap=self.cmap, extent=[0, len(array), len(array), 0], norm=norm)
self.colorbar = None
self['root'].autoscale(False)
self['root'].get_xaxis().set_ticks([])
self['root'].get_yaxis().set_ticks([])
[docs] def add_colorbar(self, loc='right', size='10%', pad=0.1,
new_ax_name='colorbar'):
"""
Adds a colorbar to the heatmap in a new axis.
Parameters
----------
loc : {'top', 'bottom', 'left', 'right'}
Which side of the heatmap to add the new colorbar to.
size : str
The size of the new axis as a percentage of the main heatmap width.
Should be passed as a string ending in '%'.
pad : float
The padding to use between the existing parts of the figure and the
newly added axis.
new_ax_name : str
A name for the new axis. The axis will be accessible as
``h[new_ax_name]`` where ``h`` is this ExtendableHeatmap instance.
Returns
-------
pyplot colorbar
The newly added colorbar.
"""
self.add_ax(new_ax_name, loc, size, pad)
self.colorbar = self.fig.colorbar(self.im, cax=self.axes[new_ax_name])
if self.categorical:
n = len(self.color_keys)
self.colorbar.set_ticks((np.arange(n) + 0.5) * (n - 1) / float(n))
self.colorbar.set_ticklabels(self.color_keys)
return self.colorbar
[docs] def add_margin_ax(self, loc='bottom', size='10%', pad=0.0,
new_ax_name='new_h_axis', axis_limits=(0, 1)):
"""
Adds a new axis to the margin of this ExtendableHeatmap.
Parameters
----------
loc : {'top', 'bottom', 'left', 'right'}
Which side of the figure to add the new axis to.
size : str
The size of the new axis as a percentage of the main heatmap width.
Should be passed as a string ending in '%'.
pad : float
The padding to use between the existing parts of the figure and the
newly added axis.
new_ax_name : str
A name for the new axis. The axis will be accessible as
``h[new_ax_name]`` where ``h`` is this ExtendableHeatmap instance.
Returns
-------
pyplot axis
The newly added axis.
"""
# create new axis
ax = self.add_ax(new_ax_name, loc, size, pad)
# force limits
if loc in ['bottom', 'top']:
ax.set_xlim([self.grange_x['start'], self.grange_x['end']])
if axis_limits is not None:
ax.set_ylim(axis_limits)
else:
ax.set_ylim([self.grange_y['end'], self.grange_y['start']])
if axis_limits is not None:
ax.set_xlim(axis_limits)
# clean ticks or disable axis
# ax.get_xaxis().set_ticks([])
# ax.get_yaxis().set_ticks([])
ax.axis('off')
return ax