heatmap test

app.py

@@ -6,7 +6,9 @@ import scglue
 import os
 import gzip
 import shutil
+import cooler
 import numpy as np
+import plotply.express as px
 import matplotlib.pyplot as plt
 
 from huggingface_hub import hf_hub_download
@@ -22,6 +24,7 @@ dist_range = 10e6
 window_size = 64
 eps = 1e-4
 n_neighbors = 5
+heatmap_range = 3  # 3X the distance of the interaction
 
 
 def unzip(file):
@@ -38,6 +41,9 @@ graph_file = hf_hub_download(repo_id="dylanplummer/islet-epigenome", filename=f"
 rna_file = unzip(rna_file)
 hic_file = unzip(hic_file)
 
+beta_cooler_file = hf_hub_download(repo_id="dylanplummer/islet-epigenome", filename=f"pseudobulk_coolers/Beta_50kb.cool", repo_type="dataset", token=os.environ['DATASET_SECRET'])
+beta_cooler = cooler.Cooler(beta_cooler_file)
+
 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)
@@ -105,6 +111,35 @@ def get_chromosome_from_filename(filename):
         return filename[chr_index:]
     else:
         return filename[chr_index: file_ending_index]
+    
+def get_heatmap(gene, locus1, locus2):
+    try:
+        loc = rna.var.loc[gene]
+    except KeyError:
+        print('Could not find loci', gene)
+        return None
+    chrom = loc["chrom"]
+    chromStart = loc["chromStart"]
+    chromEnd = loc["chromEnd"]
+    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)
+    heatmap_locus = f'{chrom}:{chromStart - interaction_dist * heatmap_range}:{chromEnd + interaction_dist * heatmap_range}'
+    mat = beta_cooler.matrix().fetch(heatmap_locus)
+    fig = px.imshow(mat)
+    return fig
 
 def perturb(gene, locus1, locus2):
     locus1 = locus1.replace(',', '')
@@ -233,6 +268,9 @@ with gr.Blocks(theme='WeixuanYuan/Soft_dark') as demo:
             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)
     with gr.Row():
+        out_heatmap = gr.Plot(elem_id='out-heatmap', scale=1)
+    with gr.Row():
+        heatmap_button = gr.Button(label="Generate Heatmaps", elem_id='heatmap-button', scale=1)
         run_button = gr.Button(label="Run", elem_id='run-button', scale=1)
     with gr.Row():
         out_img = gr.Image(elem_id='out-img', scale=1)
@@ -256,6 +294,7 @@ with gr.Blocks(theme='WeixuanYuan/Soft_dark') as demo:
                 outputs=outputs,
                 fn=perturb, cache_examples=os.getenv('SYSTEM') == 'spaces')
     run_button.click(perturb, inputs, outputs=outputs)
+    heatmap_button.click(get_heatmap, inputs, outputs=[out_heatmap])
 
 
 if __name__ == "__main__":

requirements.txt

@@ -4,4 +4,6 @@ scanpy
 scglue
 numpy 
 pandas<2.0.0
-matplotlib
+matplotlib
+cooler
+plotly