add inference with gnina-torch

app.py

@@ -2,10 +2,12 @@ import gradio as gr
 
 import os
 
+
 def load_html(html_file: str):
     with open(os.path.join("html", html_file), "r") as f:
         return f.read()
 
+
 def load_protein_from_file(protein_file) -> str:
     """
     Parameters
@@ -21,21 +23,22 @@ def load_protein_from_file(protein_file) -> str:
     with open(protein_file.name, "r") as f:
         return f.read()
 
+
 def load_ligand_from_file(ligand_file):
     with open(ligand_file.name, "r") as f:
         return f.read()
-    
+
+
 def protein_html_from_file(protein_file):
     protein = load_protein_from_file(protein_file)
     protein_html = load_html("protein.html")
 
     html = protein_html.replace("%%%PDB%%%", protein)
 
-    return f"""<iframe style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera; 
-    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
-    allow-scripts allow-same-origin allow-popups 
-    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
-    allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""
+    wrapper = load_html("wrapper.html")
+
+    return wrapper.replace("%%%HTML%%%", html)
+
 
 def ligand_html_from_file(ligand_file):
     ligand = load_ligand_from_file(ligand_file)
@@ -43,11 +46,10 @@ def ligand_html_from_file(ligand_file):
 
     html = ligand_html.replace("%%%SDF%%%", ligand)
 
-    return f"""<iframe style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera; 
-    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
-    allow-scripts allow-same-origin allow-popups 
-    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
-    allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""
+    wrapper = load_html("wrapper.html")
+
+    return wrapper.replace("%%%HTML%%%", html)
+
 
 def protein_ligand_html_from_file(protein_file, ligand_file):
     protein = load_protein_from_file(protein_file)
@@ -57,11 +59,70 @@ def protein_ligand_html_from_file(protein_file, ligand_file):
     html = protein_ligand_html.replace("%%%PDB%%%", protein)
     html = html.replace("%%%SDF%%%", ligand)
 
-    return f"""<iframe style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera; 
-    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
-    allow-scripts allow-same-origin allow-popups 
-    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
-    allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""
+    wrapper = load_html("wrapper.html")
+
+    return wrapper.replace("%%%HTML%%%", html)
+
+
+def predict(protein_file, ligand_file, cnn="default"):
+    import molgrid
+    from gninatorch import gnina, dataloaders
+    import torch
+    import pandas as pd
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(device)
+
+    model, ensemble = gnina.setup_gnina_model(cnn, 23.5, 0.5)
+    model.eval()
+    model.to(device)
+
+    example_provider = molgrid.ExampleProvider(
+        data_root="",
+        balanced=False,
+        shuffle=False,
+        default_batch_size=1,
+        iteration_scheme=molgrid.IterationScheme.SmallEpoch,
+    )
+
+    with open("data.in", "w") as f:
+        f.write(protein_file.name)
+        f.write(" ")
+        f.write(ligand_file.name)
+
+    print("Populating example provider... ", end="")
+    example_provider.populate("data.in")
+    print("done")
+
+    grid_maker = molgrid.GridMaker(resolution=0.5, dimension=23.5)
+
+    # TODO: Allow average over different rotations
+    loader = dataloaders.GriddedExamplesLoader(
+        example_provider=example_provider,
+        grid_maker=grid_maker,
+        random_translation=0.0,  # No random translations for inference
+        random_rotation=False,  # No random rotations for inference
+        grids_only=True,
+        device=device,
+    )
+
+    print("Loading and gridding data... ", end="")
+    batch = next(loader)
+    print("done")
+
+    print("Predicting... ", end="")
+    with torch.no_grad():
+        log_pose, affinity, affinity_var = model(batch)
+    print("done")
+
+    return pd.DataFrame(
+        {
+            "CNNscore": [torch.exp(log_pose[:, -1]).item()],
+            "CNNaffinity": [affinity.item()],
+            "CNNvariance": [affinity_var.item()],
+        }
+    )
+
 
 demo = gr.Blocks()
 
@@ -82,12 +143,21 @@ with demo:
             ligand = gr.HTML()
             lbtn.click(fn=ligand_html_from_file, inputs=[lfile], outputs=ligand)
 
+    gr.Markdown("# Protein-Ligand Complex")
     with gr.Row():
-        gr.Markdown("# Protein-Ligand Complex")
         plcomplex = gr.HTML()
 
         # TODO: Automatically display complex when both files are uploaded
         plbtn = gr.Button("View")
-        plbtn.click(fn=protein_ligand_html_from_file, inputs=[pfile, lfile], outputs=plcomplex)
+        plbtn.click(
+            fn=protein_ligand_html_from_file, inputs=[pfile, lfile], outputs=plcomplex
+        )
+
+    gr.Markdown("# Gnina-Torch")
+    with gr.Row():
+        df = gr.Dataframe()
+        btn = gr.Button("Score!")
+        btn.click(fn=predict, inputs=[pfile, lfile], outputs=df)
+
 
 demo.launch()

html/ligand.html

@@ -26,7 +26,7 @@
             let config = { backgroundColor: "white" };
             let viewer = $3Dmol.createViewer(element, config);
             viewer.addModel(sdf, "sdf");
-            viewer.getModel(0).setStyle({}, { "stick": { "colorscheme": "lightgreyCarbon" } });
+            viewer.getModel(0).setStyle({}, { "stick": { "colorscheme": "purpleCarbon" } });
             viewer.zoomTo();
             viewer.render();
             viewer.zoom(0.8, 2000);

html/pl.html

@@ -29,7 +29,7 @@
             viewer.addModel(pdb, "pdb");
             viewer.addModel(sdf, "sdf");
             viewer.getModel(0).setStyle({}, { cartoon: { colorscheme: "whiteCarbon" } });
-            viewer.getModel(1).setStyle({}, { stick: { colorscheme: "lightgreyCarbon" } });
+            viewer.getModel(1).setStyle({}, { stick: { colorscheme: "purpleCarbon" } });
             viewer.zoomTo();
             viewer.render();
             viewer.zoom(0.8, 2000);

html/wrapper.html

@@ -0,0 +1,10 @@
+<iframe style="width: 100%; height: 600px"
+    name="result"
+    allow="midi; geolocation; microphone; camera; display-capture; encrypted-media;"
+    sandbox="allow-modals allow-forms allow-scripts allow-same-origin allow-popups allow-top-navigation-by-user-activation allow-downloads"
+    allowfullscreen="" 
+    allowpaymentrequest="" 
+    frameborder="0"
+    srcdoc='%%%HTML%%%'
+    >
+</iframe>

requirements.txt

@@ -1 +1,2 @@
-gninatorch==0.0.2
+gninatorch==0.0.2
+pandas>=1.0