app.py

import gradio as gr
import anndata as ad
import networkx as nx
import scanpy as sc
import scglue
import os
import gzip
import shutil
import cooler
import matplotlib
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import PowerNorm

from huggingface_hub import hf_hub_download


prior_name = 'dcq'
resolution = '10kb'
loop_q = 0.99
suffix = f'ice_{loop_q}'
out_dir = 'pooled_cells_' + suffix
plot_dir = os.path.join(out_dir, 'ism_loop_plots_distal')
dist_range = 10e6
window_size = 64
eps = 1e-4
n_neighbors = 5


def unzip(file):
    with gzip.open(file, 'rb') as f_in:
        new_file = file.replace('.gz', '')
        with open(new_file, 'wb') as f_out:
            shutil.copyfileobj(f_in, f_out)
    return new_file


rna_file = hf_hub_download(repo_id="dylanplummer/islet-epigenome", filename="rna.h5ad.gz", repo_type="dataset", token=os.environ['DATASET_SECRET'])
hic_file = hf_hub_download(repo_id="dylanplummer/islet-epigenome", filename="hic.h5ad.gz", repo_type="dataset", token=os.environ['DATASET_SECRET'])
graph_file = hf_hub_download(repo_id="dylanplummer/islet-epigenome", filename=f"{prior_name}_prior_{resolution}_{suffix}.graphml.gz", repo_type="dataset", token=os.environ['DATASET_SECRET'])
rna_file = unzip(rna_file)
hic_file = unzip(hic_file)

cooler_celltypes = ['Alpha', 'Beta', 'Acinar', 'Duct', 'PSC']
cooler_resolutions = ['50kb', '200kb', '500kb']
coolers = {}
for res in cooler_resolutions:
    coolers[res] = []
for celltype in cooler_celltypes:
    for res in cooler_resolutions:
        cooler_file = hf_hub_download(repo_id="dylanplummer/islet-epigenome", filename=f"pseudobulk_coolers/{celltype}_{res}.cool", repo_type="dataset", token=os.environ['DATASET_SECRET'])
        coolers[res].append(cooler.Cooler(cooler_file))
beta_cooler = coolers[cooler_resolutions[0]][1]

model_file = hf_hub_download(repo_id="dylanplummer/HiGLUE-islet", filename=f"glue_hic_{prior_name}_prior_{resolution}_{suffix}.dill", repo_type="model", token=os.environ['DATASET_SECRET'])

rna = ad.read_h5ad(rna_file)
# sc.pp.neighbors(rna, metric="cosine", n_neighbors=n_neighbors)
# sc.tl.umap(rna)
scglue.models.configure_dataset(rna, "NB", use_highly_variable=True, use_rep="X_pca", use_layer="counts")

celltypes = sorted(rna.obs['celltype'].unique())
n_clusters = len(celltypes)
colors = list(plt.cm.tab10(np.int32(np.linspace(0, n_clusters + 0.99, n_clusters))))
color_map = {celltype: colors[i] for i, celltype in enumerate(celltypes)}
sorted_color_map = {celltype + '_sorted': colors[i] for i, celltype in enumerate(celltypes)}
rna_color_map = {celltype + '_rna': colors[i] for i, celltype in enumerate(celltypes)}
color_map = {**color_map, **sorted_color_map, **rna_color_map}
color_map['Other'] = 'gray'

hic = ad.read_h5ad(hic_file)
hic.var["highly_variable"] = hic.var[f"{prior_name}_highly_variable"]
prior = nx.read_graphml(graph_file)
glue = scglue.models.load_model(model_file)


genes = []
loops = []
for e, attr in dict(prior.edges).items():
    if attr["type"] == 'overlap':
        gene_name = e[0]
        if gene_name.startswith('chr'):
            gene_name = e[1]
        if gene_name not in genes and not gene_name.startswith('chr'):
            genes.append(gene_name)
    elif attr["type"] == 'hic':
        loops.append(e)
rna.var["highly_variable"] = rna.var["highly_variable"] & rna.var["in_hic"]
genes = rna.var.query(f"highly_variable").index
gene_idx_map = {}
for i, gene in enumerate(genes):
    gene_idx_map[gene] = i
peaks = hic.var.query("highly_variable").copy()

rna_recon = glue.decode_data("rna", "rna", rna, prior)


def get_closest_peak_to_gene(gene_name, rna, peaks):
    try:
        loc = rna.var.loc[gene_name]
    except KeyError:
        print('Could not find loci', gene_name)
        return None
    chrom = loc["chrom"]
    chromStart = loc["chromStart"]
    peaks['in_chr'] = peaks['chrom'] == chrom
    peaks['dist'] = peaks['chromStart'].apply(lambda s: abs(s - chromStart))
    peaks.loc[~peaks['in_chr'], 'dist'] = 1e9  # set distance to 1e9 if not in same chromosome
    return peaks['dist'].idxmin()


def get_closest_gene_to_peak(peak_name, rna, peaks):
    try:
        loc = peaks.loc[peak_name]
    except KeyError:
        print('Could not find peak', peak_name)
        return None
    chrom = loc["chrom"]
    chromStart = loc["chromStart"]
    rna.var['in_chr'] = rna.var['chrom'] == chrom
    rna.var['dist'] = rna.var['chromStart'].apply(lambda s: abs(s - chromStart))
    rna.var.loc[~rna.var['in_chr'], 'dist'] = 1e9  # set distance to 1e9 if not in same chromosome
    return rna.var['dist'].idxmin()


def get_chromosome_from_filename(filename):
    chr_index = filename.find('chr')  # index of chromosome name
    if chr_index == 0:  # if chromosome name is file prefix
        return filename[:filename.find('.')]
    file_ending_index = filename.rfind('.')  # index of file ending
    if chr_index > file_ending_index:  # if chromosome name is file ending
        return filename[chr_index:]
    else:
        return filename[chr_index: file_ending_index]
    

def draw_heatmap(matrix, color_scale, ax=None, min_val=1.001, return_image=False, return_plt_im=True):
    if color_scale != 0:
        color_scale = min(color_scale, np.max(matrix))
        breaks = np.append(np.arange(min_val, color_scale, (color_scale - min_val) / 18), np.max(matrix))
    elif np.max(matrix) < 2:
        breaks = np.arange(min_val, np.max(matrix), (np.max(matrix) - min_val) / 19)
    else:
        step = (np.quantile(matrix, q=0.98) - 1) / 18
        up = np.quantile(matrix, q=0.98) + 0.011
        if up < 2:
            up = 2
            step = 0.999 / 18
        breaks = np.append(np.arange(min_val, up, step), np.max(matrix) + 0.01)
    n_bin = 20  # Discretizes the interpolation into bins
    colors = ["#FFFFFF", "#FFE4E4", "#FFD7D7", "#FFC9C9", "#FFBCBC", "#FFAEAE", "#FFA1A1", "#FF9494", "#FF8686",
              "#FF7979", "#FF6B6B", "#FF5E5E", "#FF5151", "#FF4343", "#FF3636", "#FF2828", "#FF1B1B", "#FF0D0D",
              "#FF0000"]
    cmap_name = 'deeploop'
    # Create the colormap
    cm = matplotlib.colors.LinearSegmentedColormap.from_list(
        cmap_name, colors, N=n_bin)
    norm = matplotlib.colors.BoundaryNorm(breaks, 20)
    # Fewer bins will result in "coarser" colomap interpolation
    if ax is None:
        _, ax = plt.subplots()
    img = ax.imshow(matrix, cmap=cm, norm=norm, interpolation=None)
    if return_image:
        plt.close()
        return img.get_array()
    elif return_plt_im:
        return img
    

def get_heatmap_locus(gene, locus1, locus2, heatmap_range):
    locus1 = locus1.replace(',', '')
    locus2 = locus2.replace(',', '')
    try:
        loc = rna.var.loc[gene]
    except KeyError:
        print('Could not find loci', gene)
        return None
    chrom = loc["chrom"]
    if locus1.startswith('chr'):
        chrom = locus1.split(':')[0]
        pos1 = locus1.split(':')[1]
        a1 = peaks.query(f"chrom == '{chrom}'")['chromStart'].apply(lambda s: abs(s - int(pos1))).idxmin()
    else:
        a1 = get_closest_peak_to_gene(locus1, rna, peaks)
    if locus2.startswith('chr'):
        chrom = locus2.split(':')[0]
        pos2 = locus2.split(':')[1]
        a2 = peaks.query(f"chrom == '{chrom}'")['chromStart'].apply(lambda s: abs(s - int(pos2))).idxmin()
    else:
        a2 = get_closest_peak_to_gene(locus2, rna, peaks)
    a1_start = peaks.loc[a1, 'chromStart']
    a2_start = peaks.loc[a2, 'chromStart']
    interaction_dist = abs(a1_start - a2_start)
    chrom_size = beta_cooler.chromsizes[chrom]
    locus_start = max(0, a1_start - interaction_dist - heatmap_range * 1000)
    locus_end = min(chrom_size, a2_start + interaction_dist + heatmap_range * 1000)
    heatmap_locus = f'{chrom}:{locus_start}-{locus_end}'
    heatmap_size = abs(locus_start - locus_end)
    res = 'deeploop'
    if heatmap_size > 5000000 or res not in cooler_resolutions:
        res = '50kb'
    if heatmap_size > 10000000:
        res = '200kb'
    if heatmap_size > 20000000:
        res = '500kb'
    print(heatmap_locus, res)
    return heatmap_locus, res
    

def get_heatmap(celltype, gene, locus1, locus2, heatmap_range):
    heatmap_locus, res = get_heatmap_locus(gene, locus1, locus2, heatmap_range)
    c = coolers[res][cooler_celltypes.index(celltype)]
    mat = c.matrix().fetch(heatmap_locus)
    bins = c.bins().fetch(heatmap_locus)
    locus1 = locus1.replace(',', '')
    locus2 = locus2.replace(',', '')
    if locus1.startswith('chr'):
        a1_chrom = locus1.split(':')[0]
        a1_start = int(locus1.split(':')[1])
    else:
        try:
            loc = rna.var.loc[gene]
        except KeyError:
            print('Could not find loci', gene)
            return None
        a1_chrom = loc["chrom"]
        a1_start = loc["chromStart"]
    a1_idx = bins.query(f"chrom == '{a1_chrom}'")['start'].apply(lambda s: abs(s - a1_start)).argmin()
    if locus2.startswith('chr'):
        a2_chrom = locus2.split(':')[0]
        a2_start = int(locus2.split(':')[1])
    else:
        try:
            loc = rna.var.loc[gene]
        except KeyError:
            print('Could not find loci', gene)
            return None
        a2_chrom = loc["chrom"]
        a2_start = loc["chromStart"]
    a2_idx = bins.query(f"chrom == '{a2_chrom}'")['start'].apply(lambda s: abs(s - a2_start)).argmin()
    print(a1_idx, a2_idx)
    #img = draw_heatmap(mat, 0, return_image=True)
    fig, ax = plt.subplots(figsize=(4, 4))
    if res == 'deeploop':
        draw_heatmap(mat, 3.0, ax=ax, min_val=0.8)
    else:
        ax.imshow(mat, cmap='Reds', norm=PowerNorm(gamma=0.3), interpolation=None)
    
    # remove white padding
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
    plt.axis('off')
    plt.axis('image')

    w, h = fig.canvas.get_width_height()
    #heatmap_x = (a1_idx / len(bins)) * w
    #heatmap_y = (a2_idx / len(bins)) * h
    #print(heatmap_x, heatmap_y)
    heatmap_x = a1_idx
    heatmap_y = a2_idx
    ax.scatter(int(heatmap_x), int(heatmap_y), color='green', marker='2', s=150)

    # redraw the canvas
    fig = plt.gcf()
    fig.canvas.draw()

    # convert canvas to image using numpy
    img = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
    img = img.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    plt.close()
    return img

def get_heatmaps(gene, locus1, locus2, heatmap_range):
    res = []
    for celltype in cooler_celltypes:
        res.append(get_heatmap(celltype, gene, locus1, locus2, heatmap_range))
    alpha, beta, acinar, duct, psc = res
    return alpha, beta, acinar, duct, psc


def perturb(gene, locus1, locus2, use_max_gene, tf_knockout):
    if tf_knockout:
        guidance_hvf = prior.copy()
        try:
            guidance_hvf.remove_node(gene)
        except:
            pass
        # find index of tf and set all counts and values to zero
        gene_idx = gene_idx_map[gene]
        a1 = get_closest_peak_to_gene(gene, rna, peaks)
        try:
            guidance_hvf.remove_node(a1)
        except:
            pass
    else:
        locus1 = locus1.replace(',', '')
        locus2 = locus2.replace(',', '')
        res = {'feat': [], 'log2FC': [], 'var': [], 'a1_idx': [], 'a2_idx': []}
        for c in celltypes:
            res[c] = []
            res[f'{c}_var'] = []

        links_dict = {}
        for c in celltypes:
            links_dict[c] = {'chrom1': [], 'chrom1Start': [], 'chrom1End': [], 'chrom2': [], 'chrom2Start': [], 'chrom2End': [], 'score': [], 'strand1': [], 'strand2': []}

        if locus1.startswith('chr'):
            chrom = locus1.split(':')[0]
            pos1 = locus1.split(':')[1]
            a1 = peaks.query(f"chrom == '{chrom}'")['chromStart'].apply(lambda s: abs(s - int(pos1))).idxmin()
        else:
            a1 = get_closest_peak_to_gene(locus1, rna, peaks)
        if locus2.startswith('chr'):
            chrom = locus2.split(':')[0]
            pos2 = locus2.split(':')[1]
            a2 = peaks.query(f"chrom == '{chrom}'")['chromStart'].apply(lambda s: abs(s - int(pos2))).idxmin()
        else:
            a2 = get_closest_peak_to_gene(locus2, rna, peaks)
        
        print(a1)
        print(a2)

        a1_closest_gene = get_closest_gene_to_peak(a1, rna, peaks)
        a2_closest_gene = get_closest_gene_to_peak(a2, rna, peaks)
        print(a1_closest_gene, a2_closest_gene)
        
        guidance_hvf = prior.copy()
        try:
            path_to_gene_a1 = nx.shortest_path(prior, source=a1, target=gene)
            try:
                guidance_hvf.remove_edge(path_to_gene_a1[-2], path_to_gene_a1[-1])
            except:
                pass
            try:
                guidance_hvf.remove_edge(path_to_gene_a1[-1], path_to_gene_a1[-2])
            except:
                pass
        except:
            pass
        try:
            path_to_gene_a2 = nx.shortest_path(prior, source=a2, target=gene)
            try:
                guidance_hvf.remove_edge(path_to_gene_a2[-2], path_to_gene_a2[-1])
            except:
                pass
            try:
                guidance_hvf.remove_edge(path_to_gene_a2[-1], path_to_gene_a2[-2])
            except:
                pass
        except:
                pass
        try:
            guidance_hvf.remove_edge(a1, a2)
            guidance_hvf.remove_edge(a2, a1)
        except:
            pass
        try:
            guidance_hvf.remove_edge(a1, a1_closest_gene)
            guidance_hvf.remove_edge(a1_closest_gene, a1)
        except:
            pass
        try:
            guidance_hvf.remove_edge(a2, a2_closest_gene)
            guidance_hvf.remove_edge(a2_closest_gene, a2)
        except:
            pass
    
    perterbed_rna_recon = glue.decode_data("rna", "rna", rna, guidance_hvf)
    ism = np.log2((perterbed_rna_recon + eps) / (rna_recon + eps))
    if tf_knockout:
        max_gene_idxs = np.abs(np.mean(ism, axis=0)).argsort()[::-1][:10]
        for max_gene_idx in max_gene_idxs:
            print(rna.var.query(f"highly_variable").index[max_gene_idx], np.mean(ism[:, max_gene_idx]))
        max_gene_idx = max_gene_idxs[1]
    else:
        max_gene_idx = np.abs(np.mean(ism, axis=0)).argmax()
    max_gene = rna.var.query(f"highly_variable").index[max_gene_idx]
    print('Max gene:', max_gene)
    
    # get integer index of gene
    if use_max_gene:
        gene_idx = gene_idx_map[max_gene]
    else:
        gene_idx = gene_idx_map[gene]
    rna.obs['log2FC'] = ism[:, gene_idx]
    rna.obs['old_mean'] = rna_recon[:, gene_idx]
    rna.obs['new_mean'] = perterbed_rna_recon[:, gene_idx] 
    # compute new count based on log2FC
    rna.obs['new_count'] = (rna.layers['counts'][:, gene_idx] + eps) * 2 ** rna.obs['log2FC'] - eps
    
    fig = sc.pl.umap(rna, 
                     color=['celltype', 'log2FC'], 
                     color_map='Spectral', 
                     wspace=0.05, 
                     legend_loc='on data', 
                     legend_fontoutline=2,
                     frameon=False, 
                     return_fig=True)
    if tf_knockout:
        fig.suptitle(f'{max_gene if use_max_gene else gene} expression after TF knockout of {gene}')
    else:
        fig.suptitle(f'{max_gene if use_max_gene else gene} expression after removing ' + a1 + '-' + a2)
    #fig.tight_layout()
    #fig.patch.set_facecolor('none')
    #fig.patch.set_alpha(0.0)

    fig.canvas.draw()
    image_from_plot = np.frombuffer(fig.canvas.tostring_argb(), dtype=np.uint8)
    image_from_plot = image_from_plot.reshape(fig.canvas.get_width_height()[::-1] + (4,))
    # convert from argb to rgba
    image_from_plot = image_from_plot[:, :, [1, 2, 3, 0]]

    fig, ax = plt.subplots(figsize=(6, 4))
    sc.pl.violin(rna, max_gene if use_max_gene else gene, groupby='celltype', layer='counts', show=False, ax=ax, rotation=90)
    #sc.pl.violin(rna, 'old_mean', groupby='celltype', show=False, ax=ax, rotation=90)
    ax.set_title(f'Normal {max_gene if use_max_gene else gene} expression')
    fig.tight_layout()
    fig.canvas.draw()
    violin_img = np.frombuffer(fig.canvas.tostring_argb(), dtype=np.uint8)
    violin_img = violin_img.reshape(fig.canvas.get_width_height()[::-1] + (4,))
    violin_img = violin_img[:, :, [1, 2, 3, 0]]

    fig, ax = plt.subplots(figsize=(6, 4))
    sc.pl.violin(rna, 'new_mean', groupby='celltype', show=False, ax=ax, rotation=90)
    ax.set_title(f'{max_gene if use_max_gene else gene} normalized expression after perturbation')
    fig.tight_layout()
    fig.canvas.draw()
    violin_img_new = np.frombuffer(fig.canvas.tostring_argb(), dtype=np.uint8)
    violin_img_new = violin_img_new.reshape(fig.canvas.get_width_height()[::-1] + (4,))
    violin_img_new = violin_img_new[:, :, [1, 2, 3, 0]]

    fig, ax = plt.subplots(figsize=(6, 4))
    sc.pl.violin(rna, 'log2FC', groupby='celltype', show=False, ax=ax, rotation=90)
    ax.set_title(f'{max_gene if use_max_gene else gene} expression change')
    ax.hlines(0, -1, len(celltypes), linestyles='dashed')
    fig.tight_layout()
    fig.canvas.draw()
    violin_img_fc = np.frombuffer(fig.canvas.tostring_argb(), dtype=np.uint8)
    violin_img_fc = violin_img_fc.reshape(fig.canvas.get_width_height()[::-1] + (4,))
    violin_img_fc = violin_img_fc[:, :, [1, 2, 3, 0]]

    return image_from_plot, violin_img, violin_img_new, violin_img_fc


def update_on_tf_ko(tf_knockout, use_max_gene):
    if tf_knockout:
        return gr.update(visible=False), True
    else:
        return gr.update(visible=True), use_max_gene


with gr.Blocks(theme='WeixuanYuan/Soft_dark') as demo:
    with gr.Row():
        with gr.Column():
            in_locus = gr.Textbox(label="Gene/TF", elem_id='in-locus', scale=1)
            with gr.Row():
                max_gene_checkbox = gr.Checkbox(label="Max Gene", elem_id='max-gene-checkbox', scale=1, checked=False)
                tf_ko = gr.Checkbox(label="TF KO", elem_id='tf-ko', scale=1, checked=False)
    with gr.Column() as heatmap_col:
        anchor1 = gr.Textbox(label="Locus 1 (gene or genomic coords)", elem_id='anchor1', scale=1)
        anchor2 = gr.Textbox(label="Locus 2 (gene or genomic coords)", elem_id='anchor2', scale=1)
        heatmap_x = gr.Textbox(label="Heatmap X", elem_id='heatmap-x', scale=1, visible=False)
        heatmap_y = gr.Textbox(label="Heatmap Y", elem_id='heatmap-y', scale=1, visible=False)
        heatmap_size = gr.Slider(label="Heatmap Range", info="kb range aroung input gene locus to expand", minimum=50, maximum=5000, value=250)
        heatmap_button = gr.Button(value="Generate Heatmaps", elem_id='heatmap-button', scale=1)
        with gr.Row():
            out_heatmaps = []
            for celltype in cooler_celltypes:
                out_heatmaps.append(gr.Image(label=celltype, elem_id=f'out-heatmap-{celltype}', scale=1))
    with gr.Row():
        run_button = gr.Button(value="Run", elem_id='run-button', scale=1)
    with gr.Row():
        out_img = gr.Image(elem_id='out-img', scale=1)
    with gr.Row():
        out_violin = gr.Image(elem_id='out-violin', scale=1)
        out_violin_new = gr.Image(elem_id='out-violin-new', scale=1)
        out_violin_fc = gr.Image(elem_id='out-violin-fc', scale=1)
        #out_plot = gr.Plot(elem_id='out-plot', scale=1)

    inputs = [in_locus, anchor1, anchor2, max_gene_checkbox, tf_ko]
    outputs = [out_img, out_violin, out_violin_new, out_violin_fc]

    gr.Examples(examples=[['INS', 'chr11:2,289,895', 'chr11:2,298,840', False, False],
                          ['BHLHA15', '', '', True, True],
                          ['FOXA2', '', '', True, True], # https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4878272/
                        #   ['OTUD3', 'chr1:20457786', 'chr1:20230413', False],
                          ['IRX2', 'chr5:2704405', 'chr5:2164879', False, False],
                          ['TSPAN1', 'TSPAN1', 'PIK3R3', False, False],
                        #   ['IRX1', 'chr5:5397612', 'chr5:4850008', False],
                          ['GCG', 'GCG', 'FAP', False, False],
                        #   ['GCG', 'chr2:163162903', 'chr2:162852164', True],
                          ['LOXL4', 'chr10:100186215', 'chr10:99913493', False, False],
                        #   ['KRT19', 'chr17:22,220,637', 'chr17:39,591,813', False],
                        #   ['LPP', 'chr3:188,097,749', 'chr3:197,916,262', False],
                          ['MAFB', 'chr20:39431654', 'chr20:39368271', True, False],
                          ['CEL', 'chr9:135,937,365', 'chr9:135,973,107', False, False]],
                inputs=inputs,
                outputs=outputs,
                fn=perturb, cache_examples=os.getenv('SYSTEM') == 'spaces')
    run_button.click(perturb, inputs, outputs=outputs)
    heatmap_button.click(get_heatmaps, [in_locus, anchor1, anchor2, heatmap_size], outputs=out_heatmaps)
    tf_ko.change(update_on_tf_ko, [tf_ko, max_gene_checkbox], outputs=[heatmap_col, max_gene_checkbox])
    anchor1.change(get_heatmaps, [in_locus, anchor1, anchor2, heatmap_size], outputs=out_heatmaps)
    anchor2.change(get_heatmaps, [in_locus, anchor1, anchor2, heatmap_size], outputs=out_heatmaps)
    def set_loop(img, gene, locus1, locus2, heatmap_range, evt: gr.SelectData):
        h, w = img.shape[:2]
        idx = evt.index
        x, y = idx[0], idx[1]
        heatmap_locus, res = get_heatmap_locus(gene, locus1, locus2, heatmap_range)
        bins = coolers[res][0].bins().fetch(heatmap_locus)
        bins_idx_x = int(x / w * len(bins))
        bins_idx_y = int(y / h * len(bins))
        new_a1 = f"{bins.iloc[bins_idx_x]['chrom']}:{bins.iloc[bins_idx_x]['start']}"
        new_a2 = f"{bins.iloc[bins_idx_y]['chrom']}:{bins.iloc[bins_idx_y]['start']}"
        return new_a1, new_a2
    for out_heatmap in out_heatmaps:
        out_heatmap.select(set_loop, [out_heatmap, in_locus, anchor1, anchor2, heatmap_size], outputs=[anchor1, anchor2])


if __name__ == "__main__":
    demo.launch(share=False)