"""The 'heatmap' command."""
from __future__ import absolute_import, division, print_function
from builtins import zip
import collections
import logging
import numpy as np
import matplotlib as mpl
from matplotlib import pyplot as plt
from matplotlib.collections import BrokenBarHCollection
from . import plots
from .genome.rangelabel import unpack_range
[docs]def do_heatmap(cnarrs, show_range=None, do_desaturate=False):
"""Plot copy number for multiple samples as a heatmap."""
_fig, axis = plt.subplots()
set_colorbar(axis)
# List sample names on the y-axis
axis.set_yticks([i + 0.5 for i in range(len(cnarrs))])
axis.set_yticklabels([c.sample_id for c in cnarrs])
axis.set_ylim(0, len(cnarrs))
axis.invert_yaxis()
axis.set_ylabel("Samples")
axis.set_axis_bgcolor('#DDDDDD')
r_chrom, r_start, r_end = unpack_range(show_range)
if r_start is not None or r_end is not None:
logging.info("Showing log2 ratios in range %s:%d-%s",
r_chrom, r_start, r_end or '*')
elif r_chrom:
logging.info("Showing log2 ratios on chromosome %s", r_chrom)
# Closes over do_desaturate
def cna2df(cna):
"""Extract a dataframe of plotting points from a CopyNumArray."""
points = cna.data.loc[:, ["start", "end"]]
points["color"] = cna.log2.apply(plots.cvg2rgb, args=(do_desaturate,))
return points
# Group each file's probes/segments by chromosome
sample_data = [collections.defaultdict(list) for _c in cnarrs]
# Calculate the size (max endpoint value) of each chromosome
chrom_sizes = collections.OrderedDict()
for i, cnarr in enumerate(cnarrs):
if r_chrom:
subcna = cnarr.in_range(r_chrom, r_start, r_end, mode="trim")
sample_data[i][r_chrom] = cna2df(subcna)
chrom_sizes[r_chrom] = max(subcna.end.iat[-1] if subcna else 0,
chrom_sizes.get(r_chrom, 0))
else:
for chrom, subcna in cnarr.by_chromosome():
sample_data[i][chrom] = cna2df(subcna)
chrom_sizes[chrom] = max(subcna.end.iat[-1] if subcna else 0,
chrom_sizes.get(r_chrom, 0))
# Closes over axis
def plot_sample_chrom(i, sample):
"""Draw the given coordinates and colors as a horizontal series."""
xranges = [(start, end - start)
for start, end in zip(sample.start, sample.end)]
bars = BrokenBarHCollection(xranges, (i, i+1),
edgecolors="none",
facecolors=sample["color"])
axis.add_collection(bars)
if show_range:
# Lay out only the selected chromosome
# Set x-axis the chromosomal positions (in Mb), title as the selection
axis.set_xlim((r_start or 0) * plots.MB,
(r_end or chrom_sizes[r_chrom]) * plots.MB)
axis.set_title(show_range)
axis.set_xlabel("Position (Mb)")
axis.tick_params(which='both', direction='out')
axis.get_xaxis().tick_bottom()
axis.get_yaxis().tick_left()
# Plot the individual probe/segment coverages
for i, sample in enumerate(sample_data):
crow = sample[r_chrom]
crow["start"] *= plots.MB
crow["end"] *= plots.MB
plot_sample_chrom(i, crow)
else:
# Lay out chromosome dividers and x-axis labels
# (Just enough padding to avoid overlap with the divider line)
chrom_offsets = plots.plot_x_dividers(axis, chrom_sizes, 1)
# Plot the individual probe/segment coverages
for i, sample in enumerate(sample_data):
for chrom, curr_offset in chrom_offsets.items():
crow = sample[chrom]
if len(crow):
crow["start"] += curr_offset
crow["end"] += curr_offset
plot_sample_chrom(i, crow)
return axis
def set_colorbar(axis):
# Create our colormap
# ENH: refactor to use colormap to colorize the BrokenBarHCollection
# - maybe also refactor plots.cvg2rgb
cmap = mpl.colors.LinearSegmentedColormap.from_list('cnvheat',
[(0, 0, .75),
(1, 1, 1),
(.75, 0, 0)])
# Add a colorbar
norm = mpl.colors.Normalize(-1.33, 1.33)
mappable = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
mappable.set_array(np.linspace(-1.33, 1.33, 30))
cbar = plt.colorbar(mappable, ax=axis, orientation='vertical',
fraction=0.04, pad=0.03, shrink=0.6,
# label="log2",
ticks=(-1, 0, 1))
cbar.set_label("log2", labelpad=0)