<jupyter_start><jupyter_text>The Mighty Pileup<jupyter_code>%matplotlib inline
from matplotlib.gridspec import GridSpec
import matplotlib.pyplot as plt
import matplotlib as mpl
import seaborn as sns
mpl.style.use('seaborn-white')
import multiprocess as mp
import numpy as np
import pandas as pd
import h5py
import scipy.sparse as sps
import bioframe
import cooltools
import cooler
import bbi
mm9 = bioframe.fetch_chromsizes('mm9')
chromsizes = bioframe.fetch_chromsizes('mm9')
chromosomes = list(chromsizes.index)
conditions = ['WT', 'dN']
binsize = 100000

cooler_paths = {    
    'WT' : f'data/UNTR.{binsize}.cool',
    'T'  : f'data/TAM.{binsize}.cool',
    'dN' : f'data/NIPBL.{binsize}.cool',
}
long_names = {
    'WT': 'Wildtype',
    'T' : 'TAM',
    'dN': 'NipblKO',
}
pal = sns.color_palette('colorblind')
colors = {
    'WT': pal[0],
    'T' : '#333333',
    'dN': pal[2],
}

clrs = {
    cond: cooler.Cooler(cooler_paths[cond]) for cond in conditions
}<jupyter_output><empty_output><jupyter_text>Single landmark pileup<jupyter_code>from cooltools import snipping
ctcf = pd.read_table('data/CtcfCtrl.mm9__VS__InputCtrl.mm9.narrowPeak_with_motif.txt.gz')
ctcf.head()
sites = ctcf.sort_values('fc', ascending=False).iloc[:1000]
sites.head()
supports = [(chrom, 0, chromsizes[chrom]) for chrom in chromosomes]

flank = 600000
windows = snipping.make_bin_aligned_windows(
    binsize, 
    sites['chrom'], 
    (sites['start_m'] + sites['end_m'])//2,
    flank_bp=flank)
windows['strand'] = sites['strand_m']
windows = snipping.assign_regions(windows, supports)
windows = windows.dropna()

print(len(windows), 'windows, after assigning supports')
windows.head()
stacks = {}
piles = {}
for cond in conditions:
    expected = pd.read_table(f'data/{long_names[cond]}.{binsize//1000}kb.expected.cis.tsv')
    snipper = snipping.ObsExpSnipper(clrs[cond], expected)
    
    stack = snipping.pileup(windows, snipper.select, snipper.snip)
    
    # mirror reflect snippets whose feature is on the opposite strand
    mask = np.array(windows.strand == '+', dtype=bool)
    stack[:, :, mask] = stack[::-1, ::-1, mask]
    
    stacks[cond] = stack
    piles[cond] = np.nanmean(stack, axis=2)
gs = GridSpec(nrows=1, ncols=len(conditions) + 1, width_ratios=[20] * len(conditions) + [1])
plt.figure(figsize=(5 * len(conditions), 5))

opts = dict(
    vmin=-0.75,
    vmax=0.75,
    extent=[-flank//1000, flank//1000, -flank//1000, flank//1000],
    cmap='coolwarm'
)

for i, cond in enumerate(conditions):
    ax = plt.subplot(gs[i])
    img = ax.matshow(
        np.log2(piles[cond]),  
        **opts)
    ax.xaxis.tick_bottom()
    if i > 0:
        ax.yaxis.set_visible(False)
    plt.title(long_names[cond])

ax = plt.subplot(gs[len(conditions)])
plt.colorbar(img, cax=ax)<jupyter_output><empty_output><jupyter_text>Piling up paired landmarks<jupyter_code>anchor_dist = 300000
anchor_flank = 10000
sites = pd.read_table('data/ctcf-sites.paired.300kb_flank10kb.tsv')
sites = sites[(sites['strand1'] == '+') & (sites['strand2'] == '-')]
snippet_flank = 200000

windows1 = snipping.make_bin_aligned_windows(
    binsize, 
    sites['chrom1'], 
    sites['mid1'],
    flank_bp=snippet_flank)
windows1['strand'] = sites['strand1']

windows2 = snipping.make_bin_aligned_windows(
    binsize, 
    sites['chrom2'], 
    sites['mid2'],
    flank_bp=snippet_flank)
windows2['strand'] = sites['strand2']

windows = pd.merge(windows1, windows2, left_index=True, right_index=True, suffixes=('1', '2'))
windows = snipping.assign_regions(windows, supports)
windows = windows.dropna()
windows.head()
stacks = {}
piles = {}
for cond in conditions:
    expected = pd.read_table(f'data/{long_names[cond]}.{binsize//1000}kb.expected.cis.tsv')
    snipper = snipping.ObsExpSnipper(clrs[cond], expected)
    stack = snipping.pileup(windows, snipper.select, snipper.snip)
    stacks[cond] = stack
    piles[cond] = np.nanmean(stack, axis=2)
gs = plt.GridSpec(nrows=1, ncols=len(conditions) + 1, width_ratios=[20] * len(conditions) + [1])
plt.figure(figsize=(6 * len(conditions), 6))

opts = dict(
    vmin=-0.5,
    vmax=0.5,
    extent=[-flank//1000, flank//1000, -flank//1000, flank//1000],
    cmap='coolwarm'
)

for i, cond in enumerate(conditions):
    ax = plt.subplot(gs[i])
    img = ax.matshow(
        np.log2(np.nanmean(stacks[cond], axis=2)), #piles[cond]),  
        **opts)
    ax.xaxis.tick_bottom()
    if i > 0:
        ax.yaxis.set_visible(False)
    plt.title(long_names[cond])

ax = plt.subplot(gs[len(conditions)])
plt.colorbar(img, cax=ax)

plt.suptitle(f'convergent CTCF sites ({anchor_dist//1000} +/- {anchor_flank//1000})kb apart\n'
             f'Hi-C resolution = {binsize//1000}kb; # of pairs = {len(windows)}')<jupyter_output><empty_output><jupyter_text>Inspect examples<jupyter_code>from ipywidgets import interact
gs = plt.GridSpec(nrows=1, ncols=2)
plt.figure(figsize=(8, 4))

direc = 'dn'
condx = 'U'
condy = 'N'
n_examples = stacks_x[direc, condx, condy].shape[2]

@interact(i=(0, n_examples-1))
def f(i):
    
    Ax = stacks_x[direc, condx, condy][:,:,i]
    ax = plt.subplot(gs[0])
    ax.matshow(np.log2(Ax), cmap='coolwarm', extent=[-flank//binsize, flank//binsize,
                                                      flank//binsize, -flank//binsize])
    plt.title(condx)
    ax.xaxis.set_visible(False)
    ax.yaxis.set_visible(False)
    plt.axvline(0, c='g', ls='--', lw=0.5)
    plt.axhline(0, c='g', ls='--', lw=0.5)
    
    Ay = stacks_y[direc, condx, condy][:,:,i]
    ax = plt.subplot(gs[1])
    ax.matshow(np.log2(Ay), cmap='coolwarm', extent=[-flank//binsize, flank//binsize,
                                                      flank//binsize, -flank//binsize])
    plt.title(condy)
    ax.xaxis.set_visible(False)
    ax.yaxis.set_visible(False)
    plt.axvline(0, c='g', ls='--', lw=0.5)
    plt.axhline(0, c='g', ls='--', lw=0.5)

    this = windows[direc, condx, condy].iloc[i]
    name = '{}:{:,}'.format(this.chrom, (this.start + this.end)//2)
    plt.suptitle(name)<jupyter_output><empty_output><jupyter_text>Piledriver!An extremely hacky interactive pileup explorer<jupyter_code>import matplotlib as mpl
import seaborn as sns
import bokeh
import numpy as np

from bokeh.io import  output_file, output_notebook, push_notebook
from bokeh.plotting import figure, show, curdoc
from bokeh.layouts import (
    row, 
    column, 
    widgetbox, 
    gridplot,
    layout,
    Spacer,
    GridSpec,
)
from bokeh.models import (
    Range1d,
    GlyphRenderer, 
    LinearColorMapper, 
    LogColorMapper, 
    LogTicker, 
    ColorBar,
    Slider,
    RangeSlider,
    TextInput,
    Select,
    CustomJS,
    CustomJSFilter,
    CustomJSTransform
)

output_notebook()

CALLBACK_TEMPLATE = """
if (IPython.notebook.kernel !== undefined) {{
    var kernel = IPython.notebook.kernel;
    
    arg = "(" + {} + ")"
    cmd = "{}(" + arg + ")";
    kernel.execute(cmd, {{}}, {{}});
}}
"""
SCALAR_ARG = """cb_obj.value"""
TUPLE_ARG = """cb_obj.value[0] + "," + cb_obj.value[1]"""


def make_callback(func, is_tuple=False):
    arg_parsed = TUPLE_ARG if is_tuple else SCALAR_ARG
    return CustomJS(code=CALLBACK_TEMPLATE.format(arg_parsed, func.__name__))


def make_colormapper(cmap, vmin=None, vmax=None, scale='linear'):
    cmap = mpl.cm.get_cmap(cmap)
    palette = [mpl.colors.rgb2hex(rgb) for rgb in cmap(np.arange(cmap.N))]
    mapper_cls = {'linear': LinearColorMapper, 'log': LogColorMapper}[scale]
    return mapper_cls(palette=palette, low=vmin, high=vmax)


def _slider_on_change(value):
    lo, hi = value
    for name in app.snippet_names:
        subset = app.snippets[name][::-1, :, lo:hi]
        img = np.nanmean(subset, axis=2)
        app.images[name].data_source.data['image'] = [img]
    for name in app.track_names:
        app.quads[name].data_source.data['bottom'] = [lo]
        app.quads[name].data_source.data['top'] = [hi]
    push_notebook(handle=bokeh_handle)

    
class PileDriver:
    
    def __init__(self, 
                 tracks, 
                 snippets,  
                 tdims=(350, 750), tflank=None,
                 sdims=(350, 350), sflank=None,
                 cmap_options=None,
                 step=1):
        """

        Parameters
        ----------
        tracks : dict or tuple of (dict, list)
            Dictionary of name -> track stack array (2d). Optionally provide
            a list with subset of keys to use.
        snippets : dict or tuple of (dict, list)
            Dictionary of name -> snippet stack array (3d). Optionally provide
            a list with subset of keys to use.
        flank : int
            Radius to use for defining plot axis extents.
        sdims : tuple (int, int)
            Size of snippet figure in pixels
        tdims : tuple (int, int)
            Size of snippet figure in pixels
        cmap_options : dict
            Nested dict of colormap options for each plot: name -> dict()
            Options are: 
                * 'cmap': a matplotlib colormap object or name of one
                * 'vmin', 'vmax': range of values to map to
                * 'scale': how to map the colors to data ('linear' or 'log')
            Any missing options are substituted with default values.

        """
        if isinstance(tracks, tuple):
            self.tracks, self.track_names = tracks
        else:
            self.tracks, self.track_names = tracks, list(tracks.keys())

        if isinstance(snippets, tuple):
            self.snippets, self.snippet_names = snippets
        else:
            self.snippets, self.snippet_names = snippets, list(snippets.keys())
        
        x = next(iter(self.tracks.values()))
        n_windows = x.shape[0]
        if tflank is None:
            tflank = x.shape[1] // 2
        x = next(iter(self.snippets.values()))
        if sflank is None:
            sflank = x.shape[1] // 2
        
        self.figures = {}
        self.images = {}
        self.widgets = {}
        self.quads = {}
        
        self.dx = self.dy = 2 * sflank
        self.x_range = Range1d(-sflank, sflank)
        self.y_range = Range1d(-sflank, sflank)
        
        self.du =2 * tflank
        self.dv = n_windows
        self.u_range = Range1d(-tflank, tflank)
        self.v_range = Range1d(n_windows, 0)

        self._make_track_plots(tdims[0], tdims[1], cmap_options or {})
        self._make_snippet_plots(sdims[0], sdims[1], cmap_options or {})
        self._make_widgets(step)
    
    def _make_track_plots(self, plot_width, plot_height, color_opts):
        for name in self.track_names:
            self.figures[name] = figure(
                plot_width=plot_width, 
                plot_height=plot_height,
                x_range=self.u_range, 
                y_range=self.v_range,
                title=name,
                x_axis_location="above",
            )
            self.figures[name].title.align = 'center'
            cmap = color_opts.get(name, {}).get('cmap', 'Reds')
            vmin = color_opts.get(name, {}).get('vmin', 0)
            vmax = color_opts.get(name, {}).get('vmax', 2)
            scale = color_opts.get(name, {}).get('scale', 'linear')
            self.images[name] = self.figures[name].image(
                image=[ self.tracks[name][::-1, :] ],
                x=[-self.du/2],
                dw=[self.du],
                y=[self.dv],
                dh=[self.dv], 
                color_mapper=make_colormapper(cmap, vmin, vmax, scale), 
                name='stack1d_' + name
            )
            self.quads[name] = self.figures[name].quad(
                top=[0],
                bottom=[100],
                left=[-self.du/2],
                right=[self.du/2],
                color='#555555',
                alpha=0.6)
        
    def _make_snippet_plots(self, plot_width, plot_height, color_opts):
        for name in self.snippet_names:
            self.figures[name] = figure(
                plot_width=plot_width, 
                plot_height=plot_height,
                x_range=self.x_range, 
                y_range=self.y_range,
                title=name,
            )
            self.figures[name].title.align = 'center'
            cmap = color_opts.get(name, {}).get('cmap', 'coolwarm')
            vmin = color_opts.get(name, {}).get('vmin', 1/1.5)
            vmax = color_opts.get(name, {}).get('vmax', 1.5)
            scale = color_opts.get(name, {}).get('scale', 'log')
            self.images[name] = self.figures[name].image(
                image=[ self.snippets[name][::-1, :, 0] ],
                x=[-self.dx/2], 
                dw=[self.dx],
                y=[-self.dy/2],
                dh=[self.dy], 
                color_mapper=make_colormapper(cmap, vmin, vmax, scale), 
                name='stack2d_' + name
            )
            self.figures[name].line(x=[0, 0], y=[-self.dx/2, self.dx/2])
            self.figures[name].line(x=[-self.dy/2, self.dy/2], y=[0, 0])
        
    def _make_widgets(self, step):
        self.widgets['slider'] = RangeSlider(
            title="snippets", 
            start=0, 
            end=self.dv - 1, 
            step=step, 
            value=(0, 1),
            callback=make_callback(_slider_on_change, is_tuple=True))
    
    def get_layout(self):
        self.layout = layout(
            [self.figures[name] for name in self.snippet_names] + [self.widgets['slider']], 
            [self.figures[name] for name in self.track_names]
        )
        return self.layout
supports = [(chrom, 0, chromsizes[chrom]) for chrom in chromosomes]

binsize = 10000
cooler_paths = {
    'U': f'/net/levsha/share/lab/TCCSpitz2015/coolers/UNTR.{binsize//1000}kb.cool',
    'T': f'/net/levsha/share/lab/TCCSpitz2015/coolers/TAM.{binsize//1000}kb.cool',
    'N': f'/net/levsha/share/lab/TCCSpitz2015/coolers/NIPBL.{binsize//1000}kb.cool',
    'C': f'/net/levsha/share/lab/TCCSpitz2015/coolers/NCAPH2.{binsize//1000}kb.cool',
}
clrs = {
    cond: cooler.Cooler(cooler_paths[cond]) for cond in conditions
}

sflank = 250000
windows = snipping.make_bin_aligned_windows(
    binsize, 
    PE['chrom'], 
    PE['start'].where(PE['strand'] == '+', PE['end']),
    flank_bp=sflank)
windows['strand'] = PE['strand']
windows = snipping.assign_regions(windows, supports)
windows = windows.dropna()
print(len(windows), 'windows, after assigning supports')

stacks2d = {}
piles = {}
with mp.Pool(len(supports)) as pool:
    for cond in conditions:
        expected = pd.read_table(f'data/{long_names[cond]}.{binsize//1000}kb.expected.cis.tsv')
        snipper = snipping.ObsExpSnipper(clrs[cond], expected)
        stack = snipping.pileup(windows, snipper.select, snipper.snip, map=pool.map)
        mask = np.array(windows.strand == '+', dtype=bool)
        stack[:, :, mask] = stack[::-1, ::-1, mask]
        stacks2d[cond] = stack
        piles[cond] = np.nanmean(stack, axis=2)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14898328971209987, 'fraction_numerical': 0.02087108247768606, 'mean_word_length': 2.6187469645458963, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 26, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '647768', 'license': 'MIT', 'n_tokens_mistral': 5627, 'n_tokens_neox': 5204, 'n_words': 980, 'path': 'notebooks/05_analysis_cooltools-snipping-pileups.ipynb', 'repo_name': 'burakalver/hic-data-analysis-bootcamp', 'size': '14901'}