add all heatmaps

app.py

@@ -41,8 +41,12 @@ 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)
+celltypes = ['Alpha', 'Beta', 'Delta', 'PP', 'Acinar', 'Duct', 'PSC', 'Endothelial']
+coolers = []
+for celltype in celltypes:
+    cooler_file = hf_hub_download(repo_id="dylanplummer/islet-epigenome", filename=f"pseudobulk_coolers/{celltype}_50kb.cool", repo_type="dataset", token=os.environ['DATASET_SECRET'])
+    coolers.append(cooler.Cooler(cooler_file))
+beta_cooler = coolers[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'])
 
@@ -173,14 +177,15 @@ def get_heatmap_locus(gene, locus1, locus2, heatmap_range):
     a2_start = peaks.loc[a2, 'chromStart']
     interaction_dist = abs(a1_start - a2_start)
     chrom_size = beta_cooler.chromsizes[chrom]
-    heatmap_locus = f'{chrom}:{max(0, chromStart - interaction_dist - heatmap_range)}-{min(chrom_size, chromEnd + interaction_dist + heatmap_range)}'
+    heatmap_locus = f'{chrom}:{max(0, chromStart - interaction_dist - heatmap_range * 1000)}-{min(chrom_size, chromEnd + interaction_dist + heatmap_range * 1000)}'
     print(heatmap_locus)
     return heatmap_locus
     
 
-def get_heatmap(gene, locus1, locus2, heatmap_range):
+def get_heatmap(celltype, gene, locus1, locus2, heatmap_range):
     heatmap_locus = get_heatmap_locus(gene, locus1, locus2, heatmap_range)
-    mat = beta_cooler.matrix().fetch(heatmap_locus)
+    c = coolers[celltypes.index(celltype)]
+    mat = c.matrix().fetch(heatmap_locus)
     #img = draw_heatmap(mat, 0, return_image=True)
     fig, ax = plt.subplots(figsize=(4, 4))
     ax.imshow(mat, cmap='Reds', norm=PowerNorm(gamma=0.5), interpolation=None)
@@ -199,6 +204,12 @@ def get_heatmap(gene, locus1, locus2, heatmap_range):
     plt.close()
     return img
 
+def get_heatmaps(gene, locus1, locus2, heatmap_range):
+    res = []
+    for celltype in celltypes:
+        res.append(get_heatmap(celltype, gene, locus1, locus2, heatmap_range))
+    return res
+
 
 def perturb(gene, locus1, locus2):
     locus1 = locus1.replace(',', '')
@@ -328,7 +339,9 @@ with gr.Blocks(theme='WeixuanYuan/Soft_dark') as demo:
             anchor2 = gr.Textbox(label="Locus 2 (gene or genomic coords)", elem_id='anchor2', scale=1)
             heatmap_size = gr.Slider(label="Heatmap Range", info="kb range aroung input gene locus to expand", minimum=50, maximum=30000, value=2000)
     with gr.Row():
-        out_heatmap = gr.Image(elem_id='out-heatmap', scale=1)
+        out_heatmaps = []
+        for celltype, c in zip(celltypes, coolers):
+            out_heatmaps.append(gr.Image(label=c, elem_id=f'out-heatmap-{celltype}', scale=1))
     with gr.Row():
         heatmap_button = gr.Button(value="Generate Heatmaps", elem_id='heatmap-button', scale=1)
         run_button = gr.Button(value="Run", elem_id='run-button', scale=1)
@@ -354,7 +367,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, [in_locus, anchor1, anchor2, heatmap_size], outputs=[out_heatmap])
+    heatmap_button.click(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
@@ -366,7 +379,8 @@ with gr.Blocks(theme='WeixuanYuan/Soft_dark') as demo:
         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
-    out_heatmap.select(set_loop, [out_heatmap, in_locus, anchor1, anchor2, heatmap_size], outputs=[anchor1, anchor2])
+    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__":