Implement basic ddG prediction

app.py

@@ -1,7 +1,136 @@
+import shutil
+import tempfile
+from pathlib import Path
+from functools import partial
+
 import gradio as gr
+import numpy as np
+import torch
+
+from mutils.pdb import download_pdb
+from ppiref.extraction import PPIExtractor
+from ppiref.utils.ppi import PPIPath
+from ppiformer.tasks.node import DDGPPIformer
+from ppiformer.utils.api import predict_ddg
+from ppiformer.definitions import PPIFORMER_WEIGHTS_DIR
+
+
+def process_inputs(inputs, temp_dir):
+    pdb_code, pdb_path, partners, muts, muts_path = inputs
+
+    # Check inputs
+    if not pdb_code and not pdb_path:
+        raise gr.Error("PPI structure not specified.")
+    
+    if pdb_code and pdb_path:
+        gr.Warning("Both PDB code and PDB file specified. Using PDB file.")
+
+    if not partners:
+        raise gr.Error("Partners not specified.")
+
+    if not muts and not muts_path:
+        raise gr.Error("Mutations not specified.")
+    
+    if muts and muts_path:
+        gr.Warning("Both mutations and mutations file specified. Using mutations file.")
+
+    # Prepare PDB input
+    if pdb_path:
+        pdb_path = Path(pdb_path)
+    else:
+        try:
+            pdb_code = pdb_code.strip().lower()
+            pdb_path = temp_dir / f'pdb/{pdb_code}.pdb'
+            download_pdb(pdb_code, path=pdb_path)
+        except:
+            raise gr.Error("PDB download failed.")
+        
+    partners = list(map(lambda x: x.strip(), partners.split(',')))
+
+    # Extract PPI into temp dir
+    try:
+        ppi_dir = temp_dir / 'ppi'
+        extractor = PPIExtractor(out_dir=ppi_dir, nest_out_dir=True, join=True, radius=10.0)
+        extractor.extract(pdb_path, partners=partners)
+        ppi_path = PPIPath.construct(ppi_dir, pdb_path.stem, partners)
+    except:
+        raise gr.Error("PPI extraction failed.")
+
+    # Prepare mutations input
+    if muts_path:
+        muts_path = Path(muts_path)
+        muts = muts_path.read_text()
+    
+    muts = list(map(lambda x: x.strip(), muts.split(';')))
+
+    return ppi_path, muts
+
+
+def predict(models, temp_dir, *inputs):
+    # Process input
+    ppi_path, muts = process_inputs(inputs, temp_dir)
+
+    print(ppi_path, muts)
+
+    # Predict
+    ddg, attn = predict_ddg(models, ppi_path, muts, return_attn=True)
+
+    ddg = ddg.detach().numpy().tolist()
+    df = list(zip(muts, ddg))
+    
+    return df
+
+
+app = gr.Blocks()
+with app:
+
+    # Input GUI
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## PPI structure")
+            with gr.Row():
+                pdb_code = gr.Textbox(placeholder="1BUI", label="PDB code")
+                partners = gr.Textbox(placeholder="A,B,C", label="Partners")
+            pdb_path = gr.File(file_count="single", label="Or PDB file instead of PDB code")
+
+        with gr.Column():
+            gr.Markdown("## Mutations")
+            muts = gr.Textbox(placeholder="SC16A,FC47A", label="List of (multi-point) mutations")
+            muts_path = gr.File(file_count="single", label="Or file with mutations")
+
+    examples = gr.Examples(
+        examples=[["1BUI", "A,B,C", "SC16A;FC47A;SC16A,FC47A"]],
+        inputs=[pdb_code, partners, muts],
+        label="Examples (press line to fill)"
+    )
+
+    # Predict GUI
+    predict_button = gr.Button(value="Predict effects of mutations on PPI", variant="primary")
+
+    # Output GUI
+    gr.Markdown("## Predictions")
+    df = gr.Dataframe(
+        headers=["Mutation", "ddG"],
+        datatype=["str", "number"],
+        col_count=(2, "fixed"),
+    )
+
+    # Load models
+    models = [
+        DDGPPIformer.load_from_checkpoint(
+            PPIFORMER_WEIGHTS_DIR / f'ddg_regression/{i}.ckpt',
+            map_location=torch.device('cpu')
+        ).eval()
+        for i in range(3)
+    ]
+
+    # Create temporary directory for storing downloaded PDBs and extracted PPIs
+    temp_dir_obj = tempfile.TemporaryDirectory()
+    temp_dir = Path(temp_dir_obj.name)
 
-def greet(name):
-    return "Hello " + name + "!!"
+    # Main logic
+    inputs = [pdb_code, pdb_path, partners, muts, muts_path]
+    predict = partial(predict, models, temp_dir)
+    predict_button.click(predict, inputs=inputs, outputs=df)
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+app.launch()