Hi-C heatmaps.py

### do required imports
### make a dict with coolers, like this {'name':cooler.Cooler('path'), ...}
### make a dict with coordinates, smth like this: {'300kb_1': ('chr6', 80200000, 84200000), ...}


# Here is the key function to plot heatmaps like triangles
def pcolormesh_45deg(matrix_c, ax, start=0, resolution=1, *args, **kwargs):
    start_pos_vector = [start+resolution*i for i in range(len(matrix_c)+1)]
    import itertools
    n = matrix_c.shape[0]
    t = np.array([[1, 0.5], [-1, 0.5]])
    matrix_a = np.dot(np.array([(i[1], i[0])
                                for i in itertools.product(start_pos_vector[::-1],
                                                           start_pos_vector)]), t)
    x = matrix_a[:, 1].reshape(n + 1, n + 1)
    y = matrix_a[:, 0].reshape(n + 1, n + 1)
    im = ax.pcolormesh(x, y, np.flipud(matrix_c), *args, **kwargs)
    im.set_rasterized(True)
    return im


dodiff=1
conds = ['WT', 'I53A', 'KO']
comparator = 'WT'
genes = tracksClass.BedTrack({'file': '/home/s1529682/Documents/PhD/refGene_sorted_mm9.bed',
                    'title':'',
                    'labels':'off',
                    'line width':1,
                    'color':'grey'})
fosmids = tracksClass.BedTrack({'file': '/home/s1529682/Documents/PhD/Hi-C/polycomb/examples/Rob/Galaxy157-[300kb_PcG_Fosmids].bed',
                    'title':'Fosmids',
                    'label':'Fosmids',
                    'style':'flybase',
                    'line width':0,
                    'labels':'off',
                    'color':'grey'})
tracks = [[fosmids]*len(conds),
          [genes]*len(conds),
          [tracksClass.BigWigTrack({'file': '/home/s1529682/Documents/PhD/Hi-C/polycomb/ChIP-seq/ES_K27me3.bigWig',
                        'title':'WT H3K27me3',
                        'label':'WT H3K27me3',
                        'color':'#005c94ff',
                        'show data range':'no'})]*len(conds),
           [tracksClass.BigWigTrack({'file': '/home/s1529682/Documents/PhD/Hi-C/polycomb/ChIP-seq/ES_Ring1B.bigWig',
                        'title':'WT RING1B',
                        'label':'WT RING1B',
                        'color':'#005c94ff',
                        'show data range':'no'})]*len(conds)]

tracks = np.asarray(tracks)

for name, c in sorted(coords_bins.items()):
    coord = '%s:%s-%s' % c
    chrom, start, end = c

    comp_i = conds.index(comparator)
    #Two heatmaps
    mats = [np.triu(coolers[cond].matrix().fetch(coord), 0) for cond in conds]
    for m in mats:
        m[m==0] = np.nan
    maxs = [np.nanpercentile(np.diagonal(m, 1), 10) for m in mats]
    mins = [np.nanpercentile(m, 1) for m in mats]
    amin = np.min(mins)
    amax = np.max(maxs)
    f, axes = plt.subplots(nrows=1+dodiff+tracks.shape[0], ncols=len(conds),
                           figsize=(len(conds)*7, 5+5*dodiff+tracks.shape[0]/5),
                           gridspec_kw={'height_ratios':[10] + [10]*dodiff + [0.25] + [1.5]*(len(tracks)-1),
                                        'width_ratios' :[1]*len(conds),
                                        'wspace':0.05, 'hspace':0.00})
    if len(conds)==1:
        axes=axes[:, np.newaxis]
    for i, cond in enumerate(conds):
        m = mats[i]
        m[m==0] = np.nan
        m = np.ma.masked_invalid(m)
        mappable = pcolormesh_45deg(m, axes[0, i], start, coolers[cond].binsize,
                                    cmap='fall', norm=LogNorm(vmin=amin, vmax=amax))
        axes[0, i].set(ylim=0, aspect=0.5, yticks=[])
        axes[0, i].set_title(cond, size='xx-large')
            
    cax1 = inset_axes(axes[0, 0], width="7.5%", height="65%", loc='upper left')
    p = plt.colorbar(mappable, cax=cax1) 
    cax1.yaxis.tick_left()
    cax1.yaxis.set_label_position("right")
    cax1.text(0.5, 0.5, 'Normalized\ncontact frequency',
              transform=cax1.transAxes, va='center', ha='center',
              rotation=90)
    p.outline.set_visible(False)

    if dodiff:
        diffs = [mat/mats[comp_i] for mat in mats]
        mx = np.max([np.nanmax(diff) for diff in diffs])
        mn = np.min([np.nanmin(diff) for diff in diffs])
        amax = 2**np.max(np.abs(np.log2([mx, mx])))*0.5
        amin = 2**-np.log2(amax)
        diffs = [np.ma.masked_invalid(diff) for diff in diffs]
        for i, diff in enumerate(diffs):
            axes[1, i].axis('off')
            mappable = pcolormesh_45deg(diff, axes[1, i], start, coolers[conds[i]].binsize,
                                        norm=LogNorm(vmin=amin, vmax=amax), cmap='coolwarm')
            axes[1, i].set(ylim=0, aspect=0.5, yticks=[])
        cax2 = inset_axes(axes[1, -1], width="7.5%", height="65%", loc='upper right')
        p = plt.colorbar(mappable, cax=cax2,
                         format=matplotlib.ticker.LogFormatterSciNotation(base=2),
                         ticks=matplotlib.ticker.LogLocator(base=2))
        cax2.yaxis.tick_right()
        cax2.yaxis.set_label_position("left")
        cax2.text(0.5, 0.5, 'Mutant/%s' % comparator,
                  transform=cax2.transAxes, va='center', ha='center',
                  rotation=90)
        p.outline.set_visible(False)


    for (i, j), track in np.ndenumerate(tracks):
        ax = axes[i+1+dodiff, j]
        try:
            track.plot(ax, chrom, start, end)
        except KeyError:
            pass
        if 'label' in track.properties:
            ax.text(0, 0.95, track.properties['label'], va='top', ha='left', transform=ax.transAxes,
                    rotation=0, size='small', color=track.properties['color'])
        if i>0:
            ax.get_shared_y_axes().join(ax, axes[i+1+dodiff, 0])
            ax.autoscale()
        if 'bw' in track.__dict__:
            ax.set_ylim(0)

    for ax in axes.flatten():
        ax.axis('off')
    for ax in axes[:-1, :].flatten():
        ax.spines['bottom'].set_visible(False)
        ax.set_xticks([])
    for ax in axes[-2:, 0].flatten():
        ax.spines['right'].set_visible(False)
        ax.spines['left'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.axis('on')
        ax.set_yticks([ax.get_ylim()[-1]])
    for ax in axes[-1, :].flatten():
        ax.spines['right'].set_visible(False)
        ax.spines['left'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.axis('on')
        ax.xaxis.set_major_formatter(matplotlib.ticker.EngFormatter(sep=''))
    for ax in axes[:, 1:].flatten():
        ax.set_yticks([])

    axes[-1, 0].set_xlabel(chrom, ha='right')
    plt.setp(axes[-5:-3, :], yticks=[])
    ticklab = axes[-1, 0].xaxis.get_ticklabels()[0]
    trans = ticklab.get_transform()
    axes[-1, 0].xaxis.set_label_coords(axes[-1, 0].get_xlim()[0], 0, transform=trans)
    plt.setp(axes[:, :], xlim=(start, end))
    plt.savefig('/home/s1529682/Documents/PhD/Hi-C/polycomb/examples/Rob/%s_10kb.png' % name, dpi=300, bbox_inches='tight')
    plt.savefig('/home/s1529682/Documents/PhD/Hi-C/polycomb/examples/Rob/%s_10kb.pdf' % name, dpi=300, bbox_inches='tight')
    plt.close()