sparklines.py

import operator
from collections import defaultdict

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.collections import PatchCollection, LineCollection
from matplotlib.patches import Rectangle
import matplotlib.patches as mpatches


def chrom_sparklines(chrom_data, seqlens, groups=None, ranges=None,
                     col='black', alpha=1, figax=None, sharey=False, 
                     size=None, fontsize=8, bottom_spine=True, left_spine=False, 
                     hlines=None, hline_color='orange', hline_width=0.3,
                     ylabel_coord=-0.12):
    """
    groups is not implemented
    """
    nchroms = len(seqlens)
    max_len = max(seqlens.values())
    if figax is None:
        fig, ax = plt.subplots(ncols=1, nrows=nchroms, sharex=True, sharey=sharey)
    else:
        fig, ax = figax

    if groups is not None:
        unique_groups = {g: i for i, g in enumerate(sorted(set(groups)))}
        ngroups = len(unique_groups)

    # make the colormap
    if groups is not None:
        colors = mpl.cm.get_cmap(cmap)(np.arange(ngroups))

    # defaults, for no groups
    chrom_lookup = {}
    for j, (chrom, data) in enumerate(chrom_data.items()):
        chrom_lookup[chrom] = j
        fax = ax[j] # focal axes
        if groups is not None:
            group_i = unique_groups[groups[j]]
            col = colors[group_i, :]
        #fax.hlines(0, xmin = 0, xmax=seqlens[chrom], linewidth=1, color='lightgrey')
        fax.plot(data[:, 0], data[:, 1], color=col, linewidth=0.6, alpha=alpha)
        fax.set_xlim((0, max(seqlens.values())))
        fax.spines['top'].set_visible(False)
        fax.spines['right'].set_visible(False)
        if left_spine:
            fax.spines['left'].set_linewidth(0.1)
            fax.yaxis.set_tick_params(labelsize=fontsize, width=0.1, pad=1, length=3)
            #fax.spines['left'].set_bounds(0, 1000)
            #fax.set_yticks((0, 1000))
        else:
            fax.spines['left'].set_visible(False)
            fax.set_yticks([])
        
        if bottom_spine:
            fax.spines['bottom'].set_linewidth(0.1)
            fax.spines['bottom'].set_bounds(0, seqlens[chrom])
            fax.xaxis.set_tick_params(labelsize=fontsize, width=0.1)
        else:
            fax.spines['bottom'].set_visible(False)
        fax.set_ylabel(chrom, fontsize=6, rotation=0, labelpad=8, loc='bottom')
        fax.yaxis.set_label_coords(ylabel_coord, 0.5)
        fax.set_xticks([])
        
        if hlines is not None:
            lines = [[(0, h), (seqlens[chrom], h)] for h in hlines]
            lc = LineCollection(lines, colors=hline_color, linewidths=hline_width)
            fax.add_collection(lc)

        
        
    if ranges is not None:
        for (chrom, start, end) in ranges:
            ax[chrom_lookup[chrom]].hlines(0, xmin=start, xmax=end, color='red')
    if size is not None:
        fig.set_size_inches(*size)
   
    return fig, ax