add loop scatter

app.py

@@ -194,13 +194,15 @@ def get_heatmap_locus(gene, locus1, locus2, heatmap_range):
     return heatmap_locus, res
     
 
-def get_heatmap(celltype, gene, locus1, locus2, heatmap_range):
+def get_heatmap(celltype, gene, locus1, locus2, heatmap_range, heatmap_x, heatmap_y):
     heatmap_locus, res = get_heatmap_locus(gene, locus1, locus2, heatmap_range)
     c = coolers[res][cooler_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)
+    ax.imshow(mat, cmap='Reds', norm=PowerNorm(gamma=0.3), interpolation=None)
+    if heatmap_x is not None and heatmap_y is not None:
+        ax.scatter(int(heatmap_x), int(heatmap_y), color='green', marker='*', s=100)
     # remove white padding
     plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
     plt.axis('off')
@@ -216,10 +218,10 @@ def get_heatmap(celltype, gene, locus1, locus2, heatmap_range):
     plt.close()
     return img
 
-def get_heatmaps(gene, locus1, locus2, heatmap_range):
+def get_heatmaps(gene, locus1, locus2, heatmap_range, heatmap_x, heatmap_y):
     res = []
     for celltype in cooler_celltypes:
-        res.append(get_heatmap(celltype, gene, locus1, locus2, heatmap_range))
+        res.append(get_heatmap(celltype, gene, locus1, locus2, heatmap_range, heatmap_x, heatmap_y))
     alpha, beta, delta, pp, acinar, duct, psc, endothelial = res
     return alpha, beta, delta, pp, acinar, duct, psc, endothelial
 
@@ -350,10 +352,12 @@ with gr.Blocks(theme='WeixuanYuan/Soft_dark') as demo:
             in_locus = gr.Textbox(label="Target Gene", elem_id='in-locus', scale=1)
             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=30000, value=2000)
     with gr.Row():
         out_heatmaps = []
-        for celltype, c in zip(cooler_celltypes, coolers):
+        for celltype in cooler_celltypes:
             out_heatmaps.append(gr.Image(label=celltype, elem_id=f'out-heatmap-{celltype}', scale=1))
     with gr.Row():
         heatmap_button = gr.Button(value="Generate Heatmaps", elem_id='heatmap-button', scale=1)
@@ -380,7 +384,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_heatmaps, [in_locus, anchor1, anchor2, heatmap_size], outputs=out_heatmaps)
+    heatmap_button.click(get_heatmaps, [in_locus, anchor1, anchor2, heatmap_size, heatmap_x, heatmap_y], outputs=out_heatmaps)
     def set_loop(img, gene, locus1, locus2, heatmap_range, evt: gr.SelectData):
         h, w = img.shape[:2]
         idx = evt.index
@@ -393,7 +397,7 @@ with gr.Blocks(theme='WeixuanYuan/Soft_dark') as demo:
         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])
+        out_heatmap.select(set_loop, [out_heatmap, in_locus, anchor1, anchor2, heatmap_size], outputs=[anchor1, anchor2, heatmap_x, heatmap_y])
 
 
 if __name__ == "__main__":