regulatory demo

app/main.py

@@ -1,13 +1,12 @@
-#%%
 import argparse
 import os
 
 import gradio as gr
 import matplotlib.pyplot as plt
+import pandas as pd
 import pkg_resources
-from proscope.data import get_genename_to_uniprot, get_lddt, get_seq
-import pandas as pd 
 from dash_bio import Clustergram
+from proscope.data import get_genename_to_uniprot, get_lddt, get_seq
 
 seq = get_seq()
 genename_to_uniprot = get_genename_to_uniprot()
@@ -23,38 +22,46 @@ from proscope.af2 import AFPairseg
 from proscope.protein import Protein
 from proscope.viewer import view_pdb_html
 
-#%%
 args = argparse.ArgumentParser()
 args.add_argument("-p", "--port", type=int, default=7860, help="Port number")
 args.add_argument("-s", "--share", action="store_true", help="Share on network")
 args.add_argument("-d", "--data", type=str, default="/data", help="Data directory")
-# args = args.parse_args()
+args = args.parse_args()
 # set pseudo args
-args = args.parse_args(['-p', '7869', '-s', '-d', '/manitou/pmg/users/xf2217/demo_data'])
-#%%
+# args = args.parse_args(['-p', '7869', '-s', '-d', '/manitou/pmg/users/xf2217/demo_data'])
 gene_pairs = glob(f"{args.data}/structures/causal/*")
 gene_pairs = [os.path.basename(pair) for pair in gene_pairs]
-GET_CONFIG = load_config('/manitou/pmg/users/xf2217/atac_rna_data_processing/atac_rna_data_processing/config/GET')
-GET_CONFIG.celltype.jacob=True
-GET_CONFIG.celltype.num_cls=2
-GET_CONFIG.celltype.input=True
-GET_CONFIG.celltype.embed=True
-GET_CONFIG.celltype.data_dir = '/manitou/pmg/users/xf2217/pretrain_human_bingren_shendure_apr2023/fetal_adult/'
-GET_CONFIG.celltype.interpret_dir='/manitou/pmg/users/xf2217/Interpretation_all_hg38_allembed_v4_natac/'
-GET_CONFIG.motif_dir = '/manitou/pmg/users/xf2217/interpret_natac/motif-clustering'
+GET_CONFIG = load_config(
+    "/manitou/pmg/users/xf2217/atac_rna_data_processing/atac_rna_data_processing/config/GET"
+)
+GET_CONFIG.celltype.jacob = True
+GET_CONFIG.celltype.num_cls = 2
+GET_CONFIG.celltype.input = True
+GET_CONFIG.celltype.embed = True
+GET_CONFIG.celltype.data_dir = (
+    "/manitou/pmg/users/xf2217/pretrain_human_bingren_shendure_apr2023/fetal_adult/"
+)
+GET_CONFIG.celltype.interpret_dir = (
+    "/manitou/pmg/users/xf2217/Interpretation_all_hg38_allembed_v4_natac/"
+)
+GET_CONFIG.motif_dir = "/manitou/pmg/users/xf2217/interpret_natac/motif-clustering"
 motif = NrMotifV1.load_from_pickle(
     pkg_resources.resource_filename("atac_rna_data_processing", "data/NrMotifV1.pkl"),
-    GET_CONFIG.motif_dir
+    GET_CONFIG.motif_dir,
 )
-cell_type_annot = pd.read_csv(GET_CONFIG.celltype.data_dir.split('fetal_adult')[0] + 'data/cell_type_pretrain_human_bingren_shendure_apr2023.txt')
-cell_type_id_to_name = dict(zip(cell_type_annot['id'], cell_type_annot['celltype']))
-cell_type_name_to_id = dict(zip(cell_type_annot['celltype'], cell_type_annot['id']))
-avaliable_celltypes = sorted([cell_type_id_to_name[f.split('/')[-1]] for f in glob(GET_CONFIG.celltype.interpret_dir+'*')])
-#%%
-# fill this in...
-# set plot ppi to 100
-plt.rcParams['figure.dpi'] = 100
-
+cell_type_annot = pd.read_csv(
+    GET_CONFIG.celltype.data_dir.split("fetal_adult")[0]
+    + "data/cell_type_pretrain_human_bingren_shendure_apr2023.txt"
+)
+cell_type_id_to_name = dict(zip(cell_type_annot["id"], cell_type_annot["celltype"]))
+cell_type_name_to_id = dict(zip(cell_type_annot["celltype"], cell_type_annot["id"]))
+avaliable_celltypes = sorted(
+    [
+        cell_type_id_to_name[f.split("/")[-1]]
+        for f in glob(GET_CONFIG.celltype.interpret_dir + "*")
+    ]
+)
+plt.rcParams["figure.dpi"] = 100
 
 
 def visualize_AF2(tf_pair, a):
@@ -71,127 +78,226 @@ def visualize_AF2(tf_pair, a):
     fig4, ax4 = a.protein2.plot_plddt()
     fig5, ax5 = a.plot_score_heatmap()
     plt.tight_layout()
-    new_dropdown = update_dropdown(list(a.pairs_data.keys()), 'Segment pair')
+    new_dropdown = update_dropdown(list(a.pairs_data.keys()), "Segment pair")
     return fig1, fig2, fig3, fig4, fig5, new_dropdown, a
 
+
 def view_pdb(seg_pair, a):
     pdb_path = a.pairs_data[seg_pair].pdb
     return view_pdb_html(pdb_path), a, pdb_path
 
 
-
 def update_dropdown(x, label):
     return gr.Dropdown.update(choices=x, label=label)
 
+
 def load_and_plot_celltype(celltype_name, GET_CONFIG, cell):
     celltype_id = cell_type_name_to_id[celltype_name]
     cell = GETCellType(celltype_id, GET_CONFIG)
     cell.celltype_name = celltype_name
     gene_exp_fig = cell.plotly_gene_exp()
-    gene_exp_table = cell.gene_annot.groupby('gene_name')[['pred', 'obs', 'accessibility']].mean().reset_index()
-    return gene_exp_fig, gene_exp_table, cell
-    
+    return gene_exp_fig, cell
+
+
 def plot_gene_regions(cell, gene_name, plotly=True):
     return cell.plot_gene_regions(gene_name, plotly=plotly), cell
 
+
 def plot_gene_motifs(cell, gene_name, motif, overwrite=False):
     return cell.plot_gene_motifs(gene_name, motif, overwrite=overwrite)[0], cell
 
-def plot_motif_subnet(cell, motif_collection, m, type='neighbors', threshold=0.1):
-    return cell.plotly_motif_subnet(motif_collection, m, type=type, threshold=threshold), cell
+
+def plot_motif_subnet(cell, motif_collection, m, type="neighbors", threshold=0.1):
+    return (
+        cell.plotly_motif_subnet(motif_collection, m, type=type, threshold=threshold),
+        cell,
+    )
+
 
 def plot_gene_exp(cell, plotly=True):
     return cell.plotly_gene_exp(plotly=plotly), cell
 
+
 def plot_motif_corr(cell):
-    fig = Clustergram(data=cell.gene_by_motif.corr,
-                    column_labels=list(cell.gene_by_motif.corr.columns.values),
-                    row_labels=list(cell.gene_by_motif.corr.index),
-                    hidden_labels=['row', 'col'],
-                    link_method='average',
-                    display_ratio=0.1,
-                    width=600,
-                    height=400,
-                    color_map='rdbu_r',
-                    )
+    fig = Clustergram(
+        data=cell.gene_by_motif.corr,
+        column_labels=list(cell.gene_by_motif.corr.columns.values),
+        row_labels=list(cell.gene_by_motif.corr.index),
+        hidden_labels=["row", "col"],
+        link_method="average",
+        display_ratio=0.1,
+        width=600,
+        height=500,
+        color_map="rdbu_r",
+    )
     return fig, cell
 
-#%%
-# fill this in...
 
-# main
-if __name__ == '__main__':
-    with gr.Blocks(theme='sudeepshouche/minimalist') as demo:
-
-        seg_pairs = gr.State([''])
+if __name__ == "__main__":
+    with gr.Blocks(theme="sudeepshouche/minimalist") as demo:
+        seg_pairs = gr.State([""])
         af = gr.State(None)
         cell = gr.State(None)
-        
+
+        gr.Markdown(
+            """
+        # GET: A Foundation Model of Transcription Across Human Cell Types
+
+        _Transcriptional regulation, involving the complex interplay between regulatory sequences and proteins, 
+                    directs all biological processes. Computational models of transcriptions lack generalizability 
+                    to accurately extrapolate in unseen cell types and conditions. Here, we introduce GET, 
+                    an interpretable foundation model, designed to uncover deep regulatory patterns across 235 human fetal and adult cell types. 
+                    Relying exclusively on chromatin accessibility data and sequence information, GET achieves experimental-level accuracy 
+                    in predicting gene expression even in previously unseen cell types. GET showcases remarkable adaptability across new sequencing platforms and assays, 
+                    making it possible to infer regulatory activity across a broad range of cell types and conditions,
+                    and to uncover universal and cell type specific transcription factor interaction networks. 
+                    We tested its performance on prediction of chromatin regulatory activity, 
+                    inference of regulatory elements and regulators of fetal hemoglobin, 
+                    and identification of known physical interactions between transcription factors. 
+                    In particular, we show GET outperforms current models in predicting lentivirus-based massive parallel reporter assay readout with reduced input data. 
+                    In fetal erythroblast, we are able to identify distant (>1Mbps) regulatory regions that were missed by previous models. 
+                    In sum, we provide a generalizable and predictive cell type specific model for transcription together with catalogs of gene regulation and transcription factor interactions. 
+                    Benefit from this catalog, we are able to provide mechanistic understanding of previously unknown significance germline coding variants in disordered regions of PAX5, a lymphoma associated transcription factor._
+        """
+        )
+
         with gr.Row() as row:
             # Left column: Plot gene expression and gene regions
             with gr.Column():
+                gr.Markdown(
+                    """
+## Prediction performance
+This section allows the selection of cell types and provides a plot depicting the observed versus predicted gene expression levels.
+"""
+                )
                 with gr.Row() as row:
-                    celltype_name = gr.Dropdown(label='Cell Type', choices=avaliable_celltypes)  
-                    celltype_btn = gr.Button(value='Load & Plot Gene Expression')
-                gene_exp_plot = gr.Plot(label='Gene Expression Pred vs Obs')
-                gene_exp_table = gr.DataFrame(label='Gene Expression Table', max_rows=10)
-                
+                    celltype_name = gr.Dropdown(
+                        label="Cell Type", choices=avaliable_celltypes
+                    )
+                    celltype_btn = gr.Button(value="Load & Plot Gene Expression")
+                gene_exp_plot = gr.Plot(label="Gene Expression Pred vs Obs")
+
             # Right column: Plot gene motifs
             with gr.Column():
-                gene_name_for_region = gr.Textbox(label='Get important regions or motifs for gene:')
+                gr.Markdown(
+                    """
+## Cell-type specific regulatory inference
+This section allows the selection of a gene and provides plots of its cell-type specific regulatory regions and motifs.
+"""
+                )
+                gene_name_for_region = gr.Textbox(
+                    label="Get important regions or motifs for gene:"
+                )
                 with gr.Row() as row:
-                    region_plot_btn = gr.Button(value='Regions')
-                    motif_plot_btn = gr.Button(value='Motifs')
+                    region_plot_btn = gr.Button(value="Regions")
+                    motif_plot_btn = gr.Button(value="Motifs")
 
-                region_plot = gr.Plot(label='Gene Regions')
-                motif_plot = gr.Plot(label='Gene Motifs')
+                region_plot = gr.Plot(label="Gene Regions")
+                motif_plot = gr.Plot(label="Gene Motifs")
 
-        
+        gr.Markdown(
+            """
+## Motif Correlation and Causal Subnetworks
+
+Here, you can generate a heatmap to visualize motif correlations. Alternatively, you can explore the causal subnetworks related to specific motifs by selecting the motif and the type of subnetwork you are interested in, along with a effect size threshold.
+"""
+        )
         with gr.Row() as row:
             with gr.Column():
-                clustergram_btn = gr.Button(value='Plot Motif Correlation Heatmap')
-                clustergram_plot = gr.Plot(label='Motif Correlation')
+                clustergram_btn = gr.Button(value="Plot Motif Correlation Heatmap")
+                clustergram_plot = gr.Plot(label="Motif Correlation")
 
-            
             # Right column: Motif subnet plot
             with gr.Column():
                 with gr.Row() as row:
-                    motif_for_subnet = gr.Dropdown(label='Motif Causal Subnetwork', choices=motif.cluster_names)
-                    subnet_type = gr.Dropdown(label='Type', choices=['neighbors', 'parents', 'children'], default='neighbors')
+                    motif_for_subnet = gr.Dropdown(
+                        label="Motif Causal Subnetwork", choices=motif.cluster_names
+                    )
+                    subnet_type = gr.Dropdown(
+                        label="Type",
+                        choices=["neighbors", "parents", "children"],
+                        default="neighbors",
+                    )
                     # slider for threshold 0.01-0.2
-                    subnet_threshold = gr.Slider(label='Threshold', minimum=0.01, maximum=0.25, step=0.01, value=0.1)
-                subnet_btn = gr.Button(value='Plot Motif Causal Subnetwork')
-                subnet_plot = gr.Plot(label='Motif Causal Subnetwork')
-                
+                    subnet_threshold = gr.Slider(
+                        label="Threshold",
+                        minimum=0.01,
+                        maximum=0.25,
+                        step=0.01,
+                        value=0.1,
+                    )
+                subnet_btn = gr.Button(value="Plot Motif Causal Subnetwork")
+                subnet_plot = gr.Plot(label="Motif Causal Subnetwork")
 
+        gr.Markdown(
+            """
+## Structural atlas of TF-TF and TF-EP300 interactions
+
+This section allows you to explore transcription factor pairs. You can visualize various metrics such as Heatmaps and pLDDT (predicted Local Distance Difference Test) for both proteins in the interacting pair. You can also download the PDB file for specific segment pairs.
+"""
+        )
         with gr.Row() as row:
             with gr.Column():
                 with gr.Row() as row:
-                    tf_pairs = gr.Dropdown(label='TF pair', choices=gene_pairs)
-                    tf_pairs_btn = gr.Button(value='Load & Plot')
-                interact_plddt1 = gr.Plot(label='Interact pLDDT 1')
-                interact_plddt2 = gr.Plot(label='Interact pLDDT 2')
-                protein1_plddt = gr.Plot(label='Protein 1 pLDDT')
-                protein2_plddt = gr.Plot(label='Protein 2 pLDDT')
-
-                heatmap = gr.Plot(label='Heatmap')
-            
+                    tf_pairs = gr.Dropdown(label="TF pair", choices=gene_pairs)
+                    tf_pairs_btn = gr.Button(value="Load & Plot")
+                heatmap = gr.Plot(label="Heatmap")
+                interact_plddt1 = gr.Plot(label="Interact pLDDT 1")
+                interact_plddt2 = gr.Plot(label="Interact pLDDT 2")
+                protein1_plddt = gr.Plot(label="Protein 1 pLDDT")
+                protein2_plddt = gr.Plot(label="Protein 2 pLDDT")
+
             with gr.Column():
                 with gr.Row() as row:
-                    segpair = gr.Dropdown(label='Seg pair', choices=seg_pairs.value)
-                    segpair_btn = gr.Button(value='Get PDB')
+                    segpair = gr.Dropdown(label="Seg pair", choices=seg_pairs.value)
+                    segpair_btn = gr.Button(value="Get PDB")
                 pdb_html = gr.HTML(label="PDB HTML")
-                pdb_file = gr.File(label='Download PDB')
+                pdb_file = gr.File(label="Download PDB")
 
-        tf_pairs_btn.click(visualize_AF2, inputs = [tf_pairs, af], outputs = [ interact_plddt1, interact_plddt2, protein1_plddt, protein2_plddt, heatmap, segpair, af])
-        segpair_btn.click(view_pdb, inputs=[segpair, af], outputs=[pdb_html, af, pdb_file])
-        celltype_btn.click(load_and_plot_celltype, inputs=[celltype_name, gr.State(GET_CONFIG), cell], outputs=[gene_exp_plot, gene_exp_table, cell])
-        region_plot_btn.click(plot_gene_regions, inputs=[cell, gene_name_for_region], outputs=[region_plot, cell])
-        motif_plot_btn.click(plot_gene_motifs, inputs=[cell, gene_name_for_region, gr.State(motif)], outputs=[motif_plot, cell])
-        clustergram_btn.click(plot_motif_corr, inputs=[cell], outputs=[clustergram_plot, cell])
-        subnet_btn.click(plot_motif_subnet, inputs=[cell, gr.State(motif), motif_for_subnet, subnet_type, subnet_threshold], outputs=[subnet_plot, cell])
+        tf_pairs_btn.click(
+            visualize_AF2,
+            inputs=[tf_pairs, af],
+            outputs=[
+                interact_plddt1,
+                interact_plddt2,
+                protein1_plddt,
+                protein2_plddt,
+                heatmap,
+                segpair,
+                af,
+            ],
+        )
+        segpair_btn.click(
+            view_pdb, inputs=[segpair, af], outputs=[pdb_html, af, pdb_file]
+        )
+        celltype_btn.click(
+            load_and_plot_celltype,
+            inputs=[celltype_name, gr.State(GET_CONFIG), cell],
+            outputs=[gene_exp_plot, cell],
+        )
+        region_plot_btn.click(
+            plot_gene_regions,
+            inputs=[cell, gene_name_for_region],
+            outputs=[region_plot, cell],
+        )
+        motif_plot_btn.click(
+            plot_gene_motifs,
+            inputs=[cell, gene_name_for_region, gr.State(motif)],
+            outputs=[motif_plot, cell],
+        )
+        clustergram_btn.click(
+            plot_motif_corr, inputs=[cell], outputs=[clustergram_plot, cell]
+        )
+        subnet_btn.click(
+            plot_motif_subnet,
+            inputs=[
+                cell,
+                gr.State(motif),
+                motif_for_subnet,
+                subnet_type,
+                subnet_threshold,
+            ],
+            outputs=[subnet_plot, cell],
+        )
 
     demo.launch(share=args.share, server_port=args.port)
-
-
-# %%

modules/atac_rna_data_processing

@@ -1 +1 @@
-Subproject commit fc337002918de1e5f1f864e7ba94864a743fd16c
+Subproject commit aa7c0bfaac5719577e892e3226749a4d62c48848