Add UI inputs for pocket center x/y/z and flexible sidechains.

README.md

@@ -10,3 +10,15 @@ pinned: false
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+---------
+
+## How to use this space
+
+This is a simple app intended to showcase [DiffDock-Pocket](https://github.com/plainerman/DiffDock-Pocket).
+One can upload a protein and ligand, and calculate the predicted structure. The results are visualized in 3D and can be downloaded.
+
+* This app is designed to take 1 protein (in PDB format) and 1 ligand (in SDF format) at a time. For bulk inference, use the [command line interface](https://github.com/plainerman/DiffDock-Pocket/blob/main/README.md#running-diffdock-pocket-on-custom-complexes).
+
+* Our demonstration space uses a CPU, so it may take a few minutes to run. For faster results, use a GPU. 
+One can duplicate this space (at their own expense) by selecting "⋮" -> "Duplicate this space" in the top right corner, and then selecting a GPU in the "Settings" tab.

main.py

@@ -2,22 +2,25 @@ import datetime
 from typing import Tuple, Optional
 
 import gradio as gr
+import yaml
 
 import mol_viewer
 import run_utils
 
 
-def run_wrapper(protein_file, ligand_file, other_args_file, *args, **kwargs) -> Tuple[str, Optional[str], str]:
+def run_wrapper(protein_file, ligand_file, other_args_file, *args) -> Tuple[str, Optional[str], str]:
+
     if protein_file is None:
         return "Protein file is missing! Must provide a protein file in PDB format", None, ""
     if ligand_file is None:
         return "Ligand file is missing! Must provide a ligand file in SDF format", None, ""
 
+    other_args_dict = {}
     if other_args_file is not None:
-        kwargs["other_arg_file"] = other_args_file.name
+        other_args_dict = yaml.safe_load(open(other_args_file['name'], "r"))
 
     output_file, pdb_text, sdf_text = run_utils.run_cli_command(
-        protein_file.name, ligand_file.name, *args, **kwargs
+        protein_file['name'], ligand_file['name'], other_args_dict, *args,
     )
 
     output_viz = "No visualisation created"
@@ -38,15 +41,24 @@ def run():
             protein_pdb = gr.File(label="Protein PDB (required)", file_types=[".pdb"])
             ligand_sdf = gr.File(label="Ligand SDF (required)", file_types=[".sdf"])
         with gr.Row():
-            samples_per_complex = gr.Number(label="Samples Per Complex", value=1, minimum=1, maximum=100, precision=0)
+            samples_per_complex = gr.Number(label="Samples Per Complex", value=1, minimum=1, maximum=100, precision=0,
+                                            interactive=True)
             with gr.Column():
-                keep_local_structures = gr.Checkbox(label="Keep Local Structures", value=True)
-                save_vis = gr.Checkbox(label="Save Visualisation", value=True)
+                keep_local_structures = gr.Checkbox(label="Keep Local Structures", value=True, interactive=True)
+                save_vis = gr.Checkbox(label="Save Visualisation", value=True, interactive=True,
+                                       elem_id="save_visualisation")
+        with gr.Row():
+            # Add fields for pocket_center_x, pocket_center_y, pocket_center_z
+            pcx = gr.Number(label="Pocket Center X", value=None, precision=4, interactive=True)
+            pcy = gr.Number(label="Pocket Center Y", value=None, precision=4, interactive=True)
+            pcz = gr.Number(label="Pocket Center Z", value=None, precision=4, interactive=True)
+            flex_sidechains = gr.Textbox(label="Flexible Sidechains", placeholder="A:130-B:140", value=None, interactive=True,
+                                         info="Specify which amino acids will be flexible. E.g., A:130-B:140 will make amino acid with id 130 in chain A, and id 140 in chain B flexible.")
         with gr.Row():
             gr.Markdown("""Additional values can be included in "Other arguments", and will be passed 
         to [inference.py](https://github.com/plainerman/DiffDock-Pocket/blob/main/inference.py). 
         Must be a YAML file without any nesting.  """)
-            other_args = gr.File(label="Other arguments (Optional, YML)", file_types=[".yml", ".yaml"], value=None)
+            other_args_file = gr.File(label="Other arguments (Optional, YML)", file_types=[".yml", ".yaml"], value=None)
 
         with gr.Row():
             run_btn = gr.Button("Run DiffDock-Pocket")
@@ -59,9 +71,11 @@ def run():
             init_value = "Rank1 Reverse Processed Protein will be visualised here"
             viewer = gr.HTML(value=init_value, label="Protein Viewer", show_label=True)
 
-        _inputs = [protein_pdb, ligand_sdf, other_args, samples_per_complex, keep_local_structures, save_vis]
+        _inputs = [protein_pdb, ligand_sdf, other_args_file]
+        # See run_utils.py:ARG_ORDER for the order of these arguments
+        _inputs += [samples_per_complex, keep_local_structures, save_vis, pcx, pcy, pcz, flex_sidechains]
         _outputs = [message, output_file, viewer]
-        run_btn.click(fn=run_wrapper, inputs=_inputs, outputs=_outputs)
+        run_btn.click(fn=run_wrapper, inputs=_inputs, outputs=_outputs, preprocess=False)
 
     demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
 

requirements.txt

@@ -1,3 +1,3 @@
-gradio==3.50
+gradio==3.50.*
 requests==2.31.0
 pyyaml==6.0.1

run_utils.py

@@ -8,7 +8,8 @@ import uuid
 import logging
 from typing import List
 
-import yaml
+ARG_ORDER = ["samples_per_complex", "keep_local_structures", "save_visualisation",
+             "pocket_center_x", "pocket_center_y", "pocket_center_z", "flexible_sidechains"]
 
 
 def set_env_variables():
@@ -51,7 +52,7 @@ def kwargs_to_cli_args(**kwargs) -> List[str]:
             if value:
                 cli_args.append(f"--{key}")
         else:
-            if value is not None:
+            if value is not None and str(value) != "":
                 cli_args.append(f"--{key}={value}")
 
     return cli_args
@@ -67,10 +68,8 @@ def read_file_lines(fi_path: str):
 def run_cli_command(
     protein_path: str,
     ligand: str,
-    samples_per_complex: int,
-    keep_local_structures: bool,
-    save_visualisation: bool,
-    other_arg_file: str = None,
+    other_args_dict: dict,
+    *args,
     work_dir=None,
 ):
     if work_dir is None:
@@ -78,16 +77,14 @@ def run_cli_command(
             "DiffDock-Pocket-Dir", os.path.join(os.environ["HOME"], "DiffDock-Pocket")
         )
 
-    all_arg_dict = {}
-    if other_arg_file:
-        all_arg_dict = yaml.safe_load(open(other_arg_file, "r"))
-        logging.debug(f"YAML dict: {all_arg_dict}")
+    assert len(args) == len(ARG_ORDER), f'Expected {len(ARG_ORDER)} arguments, got {len(args)}'
 
-    ui_args = ["protein_path", "ligand", "samples_per_complex",
-               "keep_local_structures", "save_visualisation"]
-    for ui_arg in ui_args:
-        my_str = f"all_arg_dict['{ui_arg}'] = {ui_arg}"
-        exec(my_str)
+    all_arg_dict = other_args_dict
+    all_arg_dict["protein_path"] = protein_path
+    all_arg_dict["ligand"] = ligand
+
+    for arg_name, arg_val in zip(ARG_ORDER, args):
+        all_arg_dict[arg_name] = arg_val
 
     command = [
         "python3",
@@ -163,7 +160,7 @@ def run_cli_command(
     return full_zip_path, pdb_text, sdf_text
 
 
-if __name__ == "__main__":
+def main_test():
     # Testing code
     set_env_variables()
     configure_logging()
@@ -182,3 +179,7 @@ if __name__ == "__main__":
         save_visualisation=True,
         other_arg_file=other_arg_file
     )
+
+
+if __name__ == "__main__":
+    main_test()