Support for remote proteinmpnn calls

Author: timodonnell

README.md

@@ -4,8 +4,8 @@ Common interface to protein design tools and structure predictors
 The goal is provide a Python API that wraps recent protein design and structure prediction tools. We want to
 be able to experiment with multiple approaches for same problem without rewriting our task for each tool's API.
 
-This package also provides a large docker image with the tools installed (see [docker/base](docker/base)) as well as an 
-image that extends it to include a jupyter environment for interactive work (see [docker/full](docker/full)).
+This package also provides a huge docker image with the tools installed, including
+model weights.
 
 This is all very much WIP. Contributions are welcome.
 

api.py

@@ -11,6 +11,7 @@
 from flask_restful import reqparse, abort, Api, Resource, inputs
 
 import proteopt.alphafold
+import proteopt.proteinmpnn
 import proteopt.mock_tool
 from proteopt.common import serialize, deserialize
 
@@ -38,6 +39,7 @@ def add_argument(parser, arg_name, info, append=False):
 TOOL_CLASSES = [
     proteopt.mock_tool.MockTool,
     proteopt.alphafold.AlphaFold,
+    proteopt.proteinmpnn.ProteinMPNN,
 ]
 TOOLS = dict((cls.tool_name, cls) for cls in TOOL_CLASSES)
 
@@ -157,7 +159,8 @@ def post(self, tool_name):
 if __name__ == '__main__':
     args = arg_parser.parse_args(sys.argv[1:])
 
-    tool_configs = collections.defaultdict(dict)  # tool name -> dict
+    # tool name -> dict
+    tool_configs = dict((tool_name, {}) for tool_name in TOOLS.keys())
     for (arg, (tool, parameter)) in arg_names_to_tool_configs.items():
         tool_configs[tool][parameter] = getattr(args, arg)
 

docker/Dockerfile

@@ -45,6 +45,10 @@ RUN ln -s /data/static/alphafold-params . \
     && mkdir -p ~/.cache \
     && ln -s /data/static/omegafold_ckpt ~/.cache/omegafold_ckpt
 
+# Was getting some errors from wandb without this:
+ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
+
+
 # ****************************
 FROM minimal as base
 RUN mkdir -p /data/static
@@ -217,8 +221,6 @@ RUN cd /data/static \
 # Switch to the design-env on start:
 RUN echo "conda activate design-env" >> ~/.bashrc
 
-# Was getting some errors from wandb without this:
-ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
 
 # ****************************
 # Jupyter setup for interactive usage

docker/container-files/deploy_local_server.sh

@@ -2,9 +2,9 @@ set -e
 set -x 
 
 NUM="${1:-2}"
-NUM_PER=1
+NUM_PER="${2:-1}"
 
-source ../paths.sh
+ALPHAFOLD_WEIGHTS_DIR="/software/mlfold/alphafold-data"
 
 ENDPOINTS_FILE=/tmp/PROTEOPT_ENDPOINTS.TXT
 PIDS_FILE=/tmp/PROTEOPT_ENDPOINTS.PIDS.TXT
@@ -18,7 +18,7 @@ do
 	for j in $(seq $NUM_PER)
 	do
 		CUDA_VISIBLE_DEVICES=$(expr $i - 1) python \
-			~/git/proteopt/api.py \
+			~/proteopt/api.py \
 			--debug \
 			--alphafold-data-dir "$ALPHAFOLD_WEIGHTS_DIR" \
 			--write-endpoint-to-file /tmp/proteopt_endpoint.txt &

docker/run.sh renamed to docker/run.jupyter.sh

@@ -3,7 +3,6 @@
 docker run -it \
   -v "$(realpath ~):/root/host" \
   -v "$(realpath ..):/root/proteopt" \
-  -p 8888:8888 \
    --gpus all \
    timodonnell/proteopt-complete:latest \
    bash -c "cd /root/proteopt/docker ; bash update.sh ; cd ; bash"

docker/run.shell.sh

@@ -1,7 +1,7 @@
 import tempfile
 import os
 import json
-import pickle
+from typing import Optional
 
 import numpy
 import pandas
@@ -12,20 +12,25 @@
 
 from ProteinMPNN import protein_mpnn_run
 
+from .common import args_from_function_signature
 
 class ProteinMPNN(object):
+    tool_name = "proteinmpnn"
+    config_args = {}
+    model_args = {}
+
     def __init__(self):
         pass
 
     def run(
             self,
-            structure,
-            fixed=None,
-            num=1,
-            ca_only=False,
-            sampling_temp=0.1,
-            batch_size=1,
-            verbose=False):
+            structure : prody.Atomic,
+            fixed : Optional[prody.Atomic] = None,
+            num : int = 1,
+            ca_only : bool = False,
+            sampling_temp : float = 0.1,
+            batch_size : int = 1,
+            verbose : bool = False):
 
         # Reset resnums to avoid gaps
         chains = numpy.unique(structure.getChids())
@@ -124,5 +129,19 @@ def run(
         temp_dir.cleanup()
         return result_df
 
+    def run_multiple(self, list_of_dicts, show_progress=False):
+        results = []
+
+        iterator = list_of_dicts
+        if show_progress:
+            import tqdm
+            iterator = tqdm.tqdm(list_of_dicts)
+
+        for kwargs in iterator:
+            result = self.run(**kwargs)
+            assert result is not None
+            results.append(result)
+        return results
 
+    run_args = args_from_function_signature(run)
 

proteopt/proteinmpnn.py

@@ -46,7 +46,7 @@ def test_multiple_endpoints_mock(multiple_running_server_endpoints):
 
 def test_basic_mock(running_server_endpoint):
     client = proteopt.client.Client(endpoints=[running_server_endpoint])
-    model = client.remote_model(
+    model = client.premote_model(
         proteopt.alphafold.AlphaFold,
         model_name="MOCK",
         max_length=16,

test/test_alphafold_remote.py

@@ -13,7 +13,7 @@
 ALPHAFOLD_WEIGHTS_DIR = "/software/mlfold/alphafold-data"
 
 
-def Xtest_basic():
+def test_basic():
     region1 = prody.parsePDB(
         os.path.join(DATA_DIR, "1MBN.pdb")
     ).select("protein chain A and resid 10 to 39")

test/test_proteinmpnn.py

@@ -0,0 +1,41 @@
+import warnings
+
+import numpy.testing
+
+warnings.filterwarnings("ignore")
+
+import os
+import prody
+
+import pytest
+
+import proteopt
+import proteopt.client
+import proteopt.proteinmpnn
+
+DATA_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "data")
+
+from .util import running_server_endpoint
+
+def test_basic(running_server_endpoint):
+    client = proteopt.client.Client(endpoints=[running_server_endpoint])
+    runner = client.remote_model(
+        proteopt.proteinmpnn.ProteinMPNN)
+
+    region1 = prody.parsePDB(
+        os.path.join(DATA_DIR, "1MBN.pdb")
+    ).select("protein chain A and resid 10 to 39")
+    sequence = region1.ca.getSequence()
+    fixed_region = region1.select("resid 25 to 28 or resid 35")
+
+    results = runner.run(region1, num=5, fixed=fixed_region)
+    print(results)
+
+    assert results.shape[0] == 5
+    assert list(results.seq.str.len().unique()) == [len(region1.ca)]
+
+    for _, row in results.iterrows():
+        assert row.seq[:10] != sequence[:10]
+        assert row.seq[15:19] == sequence[15:19]
+        assert row.seq[25] == sequence[25]
+