gene plotter.

geneplot.py

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import gridspec
import seaborn as sns
sns.set_style('white')


def geneplot(exons, patho_variants, population_variants=None, constraint=None,
        density=None,
        opts={'constraint_color': (0.7, 0.7, 0.7),
              'patho_variant_color': (228/255.,26/255.,28/255.),
              'exon_color': (0.8,0.8, 0.8),
              'pop_variant_color': '#4daf4a',
              'density_window': 8,
              }):
    """
    >>> _, _ = geneplot([[1234, 1298],[22222, 22349]],
    ... [(1244, 0.001), (1247, 0.1), (1255, 0.2), (22233, 0.001), (22244, 0.2)],
    ... [(1235, 0.1), (1236, 0.2), (22340, 0.1), (22341, 0.1), (22344, 0.1)],
    ... [(1234, 1244, 81), (1244, 1255, 99), (22233, 22341, 80), (22341, 22349, 10)],
    ... density=range(1244, 1290) + range(22244, 22255) + range(22261, 22269))
    >>> plt.show()
    """

    widths = [float(e[1] - e[0]) for e in exons]
    fig = plt.figure(figsize=(8, 4))
    height_ratios =  (1, 1) if density is None else (0.8, 0.8, 1)
    gs = gridspec.GridSpec(2 if density is None else 3, len(exons), width_ratios=widths,
            height_ratios=height_ratios, hspace=0.0)

    ncols = len(exons)
    ax0 = None


    for i, exon in enumerate(exons):
        if i == 0:
            ax = plt.subplot(gs[0, i])
            ax0 = ax
        else:
            ax = fig.add_subplot(gs[0, i], sharey=ax0)

        vs = [v for v in patho_variants if exon[0] <= v[0] <= exon[1]]
        pop = [v for v in population_variants if exon[0] <= v[0] <= exon[1]]
        ctr = [v for v in constraint if exon[0] <= v[0] <= exon[1]]

        xs, ys = [], []
        for s, e, height in ctr:
            xs.extend([s, e])
            ys.extend([height, height])
        ax.step(xs, ys, color=opts['constraint_color'])
        #    ax.plot([s, e], [height, height], color=opts['constraint_color'])

        axv = fig.add_subplot(gs[1 if density is None else 2, i], sharex=ax)
        markers, stemlines, baseline = axv.stem([v[0] for v in vs], -np.log10([v[1]
            for v in vs]), linefmt='-', markerfmt='o')
        plt.setp(baseline, 'linewidth', 0)
        plt.setp(stemlines, 'color', opts['patho_variant_color'], 'zorder', -1)

        plt.setp(markers, 'color', opts['patho_variant_color'], 'zorder', 1,
                'alpha', 0.9, 'markeredgecolor', '#666666', 'mew', 1,
                'markersize', 8)

        #axe = fig.add_subplot(gs[2, i], sharex=ax)
        markers, stemlines, baseline = axv.stem([v[0] for v in pop], -np.log10([v[1] for v in pop]),
                linefmt='-', markerfmt='o')
        plt.setp(baseline, 'linewidth', 0)
        plt.setp(stemlines, 'color', opts['pop_variant_color'], 'zorder', -1)
        plt.setp(markers, 'color', opts['pop_variant_color'], 'zorder', 1,
                'alpha', 0.9, 'markeredgecolor', '#666666', 'mew', 1,
                'markersize', 8)

        if i == 0:
            ax.set_ylabel('Constraint')
            axv.set_ylabel("-log10(AF)")
        else:
            plt.setp(ax.get_yticklabels(), visible=False)
            plt.setp(axv.get_yticklabels(), visible=False)
        ax.set_yticks([])
        axv.set_yticks([])
        axv.set_xticks([])

        axv.axhspan(-0.5, 0, xmin=0.001, xmax=0.999, facecolor='none', edgecolor=opts['exon_color'],
                lw=4, zorder=9)

        if not density:
            continue
        idensity = np.zeros(exon[1] - exon[0])
        for d in density:
            if d >= exon[0] and d < exon[1]:
                idensity[d-exon[0]]+=1
        idensity = np.convolve(idensity, np.ones(opts['density_window']) / float(opts['density_window']), mode='same')
        axd = fig.add_subplot(gs[1, i], sharex=ax)
        axd.plot(np.arange(exon[0], exon[1]), idensity)
        if i == 0:
            axd.set_ylabel('Variant density')
        axd.set_yticks([])
        axd.set_xticks([])

    sns.despine(left=True, bottom=True)
    #gs.tight_layout(fig, h_pad=0)
    plt.tight_layout()
    return fig, gs
    plt.show()

if __name__ == "__main__":
    import doctest
    doctest.testmod()