Add Singularity support and some fixes

Dockerfile

@@ -8,14 +8,19 @@ RUN apt-get update && apt-get install -y software-properties-common && \
 RUN alias "python"="python3.10"
 
 # Make a virtual env that we can safely install into
-
 RUN python3 -m venv /opt/venv
+
 # Enable venv
 ENV PATH="/opt/venv/bin:$PATH"
 
 RUN pip install poetry
 
 # Set the working directory to the user's home directory
-WORKDIR /home
+# WORKDIR /home
+WORKDIR /opt/RFantibody
+
+ENV PYTHONPATH="/opt/RFantibody/src:/opt/RFantibody/src/rfantibody/rfdiffusion"
+ENV DGLBACKEND="pytorch"
 
-ENTRYPOINT /bin/bash
+# ENTRYPOINT /bin/bash  # this did not activate the environment, which was setup 
+ENTRYPOINT ["/bin/bash", "-c", "source /opt/venv/bin/activate && exec bash"]

include/setup_sif.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+export HOME=/opt/RFantibody
+cd $HOME
+
+python3 -m venv venv
+source venv/bin/activate
+pip install poetry
+
+mkdir -p dgl
+cd $HOME/include/dgl
+wget -nc https://data.dgl.ai/wheels/torch-2.3/cu118/dgl-2.4.0%2Bcu118-cp310-cp310-manylinux1_x86_64.whl
+
+cd $HOME/include/USalign
+make
+
+cd $HOME
+poetry install
+
+echo "Setup successful."

pyproject.toml

@@ -9,6 +9,7 @@ authors = [
 readme = "README.md"
 packages = [
 {include = "rfantibody", from = "src"},
+{include = "rfdiffusion", from = "src/rfantibody"},
 {include = "include"},
 {include = "se3_transformer", from = "include/SE3Transformer"}
 ]
@@ -36,6 +37,7 @@ opt-einsum = "*"
 dgl = { file = "./include/dgl/dgl-2.4.0+cu118-cp310-cp310-manylinux1_x86_64.whl" }
 e3nn = "*"
 pyrsistent = "*"
+biotite = "*"
 
 [tool.poetry.group.test]
 optional = true

scripts/examples/proteinmpnn/ab_pdb_example.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
 
-poetry run python /home/scripts/proteinmpnn_interface_design.py \
-    -pdbdir /home/scripts/examples/proteinmpnn/example_inputs \
-    -outpdbdir /home/scripts/examples/proteinmpnn/example_outputs
+poetry run python /opt/RFantibody/scripts/proteinmpnn_interface_design.py \
+    -pdbdir /opt/RFantibody/scripts/examples/proteinmpnn/example_inputs \
+    -outpdbdir /opt/RFantibody/scripts/examples/proteinmpnn/example_outputs

scripts/examples/rf2/ab_pdb_example.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
-mkdir -p /home/scripts/examples/rf2/example_outputs
+mkdir -p /opt/RFantibody/scripts/examples/rf2/example_outputs
 
-poetry run python /home/scripts/rf2_predict.py \
-    input.pdb_dir=/home/scripts/examples/rf2/example_inputs \
-    output.pdb_dir=/home/scripts/examples/rf2/example_outputs
+poetry run python /opt/RFantibody/scripts/rf2_predict.py \
+    input.pdb_dir=/opt/RFantibody/scripts/examples/rf2/example_inputs \
+    output.pdb_dir=/opt/RFantibody/scripts/examples/rf2/example_outputs

scripts/examples/rfdiffusion/antibody_pdbdesign.sh

@@ -1,14 +1,14 @@
 #!/bin/bash
 
-poetry run python  /home/src/rfantibody/scripts/rfdiffusion_inference.py \
+poetry run python  /opt/RFantibody/src/rfantibody/scripts/rfdiffusion_inference.py \
     --config-name antibody \
-    antibody.target_pdb=/home/scripts/examples/example_inputs/rsv_site3.pdb \
-    antibody.framework_pdb=/home/scripts/examples/example_inputs/hu-4D5-8_Fv.pdb \
-    inference.ckpt_override_path=/home/weights/RFdiffusion_Ab.pt \
+    antibody.target_pdb=/opt/RFantibody/scripts/examples/example_inputs/rsv_site3.pdb \
+    antibody.framework_pdb=/opt/RFantibody/scripts/examples/example_inputs/hu-4D5-8_Fv.pdb \
+    inference.ckpt_override_path=/opt/RFantibody/weights/RFdiffusion_Ab.pt \
     'ppi.hotspot_res=[T305,T456]' \
     'antibody.design_loops=[L1:8-13,L2:7,L3:9-11,H1:7,H2:6,H3:5-13]' \
     inference.num_designs=2 \
     inference.final_step=48 \
     diffuser.T=50 \
     inference.deterministic=True \
-    inference.output_prefix=/home/scripts/examples/example_outputs/ab_des
+    inference.output_prefix=/opt/RFantibody/scripts/examples/example_outputs/ab_des

scripts/examples/rfdiffusion/antibody_qvdesign.sh

@@ -1,14 +1,14 @@
 #!/bin/bash
 
-poetry run python  /home/src/rfantibody/rfdiffusion/run_inference.py \
+poetry run python  /opt/RFantibody/src/rfantibody/rfdiffusion/run_inference.py \
     --config-name antibody \
-    antibody.target_pdb=/home/scripts/examples/example_inputs/rsv_site3.pdb \
-    antibody.framework_pdb=/home/scripts/examples/example_inputs/hu-4D5-8_Fv.pdb \
-    inference.ckpt_override_path=/home/weights/RFdiffusion_Ab.pt  \
+    antibody.target_pdb=/opt/RFantibody/scripts/examples/example_inputs/rsv_site3.pdb \
+    antibody.framework_pdb=/opt/RFantibody/scripts/examples/example_inputs/hu-4D5-8_Fv.pdb \
+    inference.ckpt_override_path=/opt/RFantibody/weights/RFdiffusion_Ab.pt  \
     'ppi.hotspot_res=[T305,T456]' \
     'antibody.design_loops=[L1:8-13,L2:7,L3:9-11,H1:7,H2:6,H3:5-13]' \
     inference.num_designs=2 \
     inference.final_step=48 \
     diffuser.T=50 \
     inference.deterministic=True \
-    inference.quiver=/home/scripts/examples/example_outputs/ab_designs.qv
+    inference.quiver=/opt/RFantibody/scripts/examples/example_outputs/ab_designs.qv

scripts/examples/rfdiffusion/nanobody_pdbdesign.sh

@@ -1,14 +1,14 @@
 #!/bin/bash
 
-poetry run python  /home/src/rfantibody/scripts/rfdiffusion_inference.py \
+poetry run python  /opt/RFantibody/src/rfantibody/scripts/rfdiffusion_inference.py \
     --config-name antibody \
-    antibody.target_pdb=/home/scripts/examples/example_inputs/rsv_site3.pdb \
-    antibody.framework_pdb=/home/scripts/examples/example_inputs/h-NbBCII10.pdb \
-    inference.ckpt_override_path=/home/weights/RFdiffusion_Ab.pt \
+    antibody.target_pdb=/opt/RFantibody/scripts/examples/example_inputs/rsv_site3.pdb \
+    antibody.framework_pdb=/opt/RFantibody/scripts/examples/example_inputs/h-NbBCII10.pdb \
+    inference.ckpt_override_path=/opt/RFantibody/weights/RFdiffusion_Ab.pt \
     'ppi.hotspot_res=[T305,T456]' \
     'antibody.design_loops=[L1:8-13,L2:7,L3:9-11,H1:7,H2:6,H3:5-13]' \
     inference.num_designs=2 \
     inference.final_step=48 \
     inference.deterministic=True \
     diffuser.T=50 \
-    inference.output_prefix=/home/scripts/examples/example_outputs/nb_des
+    inference.output_prefix=/opt/RFantibody/scripts/examples/example_outputs/nb_des

scripts/examples/rfdiffusion/nanobody_qvdesign.sh

@@ -1,10 +1,10 @@
 #!/bin/bash
 
-poetry run python  /home/src/rfantibody/rfdiffusion/run_inference.py \
+poetry run python  /opt/RFantibody/src/rfantibody/rfdiffusion/run_inference.py \
     --config-name antibody \
-    antibody.target_pdb=/home/scripts/examples/example_inputs/rsv_site3.pdb \
-    antibody.framework_pdb=/home/scripts/examples/example_inputs/h-NbBCII10.pdb \
-    inference.ckpt_override_path=/home/weights/RFdiffusion_Ab.pt \
+    antibody.target_pdb=/opt/RFantibody/scripts/examples/example_inputs/rsv_site3.pdb \
+    antibody.framework_pdb=/opt/RFantibody/scripts/examples/example_inputs/h-NbBCII10.pdb \
+    inference.ckpt_override_path=/opt/RFantibody/weights/RFdiffusion_Ab.pt \
     'ppi.hotspot_res=[T305,T456]' \
     'antibody.design_loops=[L1:8-13,L2:7,L3:9-11,H1:7,H2:6,H3:5-13]' \
     inference.num_designs=2 \

scripts/proteinmpnn_interface_design.py

@@ -5,6 +5,7 @@
 import time
 
 import torch
+import numpy as np
 
 import rfantibody.proteinmpnn.util_protein_mpnn as mpnn_util
 from rfantibody.proteinmpnn.struct_manager import StructManager
@@ -40,7 +41,7 @@
                     help="The number of sequences to generate for each structure (default: 1)")
 
 # ProteinMPNN Specific Arguments
-default_ckpt = os.path.join( os.path.dirname(__file__), '/home/weights/ProteinMPNN_v48_noise_0.2.pt')
+default_ckpt = os.path.join( os.path.dirname(__file__), '/opt/RFantibody/weights/ProteinMPNN_v48_noise_0.2.pt')
 parser.add_argument("-checkpoint_path", type=str, default=default_ckpt)
 parser.add_argument("-temperature", type=float, default=0.000001, help='An a3m file containing the MSA of your target')
 parser.add_argument("-augment_eps", type=float, default=0,
@@ -53,17 +54,24 @@
 parser.add_argument("-num_connections", type=int, default=48,
                     help='Number of neighbors each residue is connected to, default 48, higher number leads to ' + \
                          'better interface design but will cost more to run the model.')
+parser.add_argument("-name_tag", type=str, default="dldesign",
+                    help=f'To append after the input file name, default: dldesign')
+parser.add_argument("-seed",  default="",
+                    help='Seed, should be in the 0-999 range, default: random seed')
 
 args = parser.parse_args(sys.argv[1:])
 
+
+
 class ProteinMPNN_runner():
     '''
     This class is designed to run the ProteinMPNN model on a single input. This class handles the loading of the model,
     the loading of the input data, the running of the model, and the processing of the output
     '''
-
+    
     def __init__(self, args, struct_manager):
         self.struct_manager = struct_manager
+        self.outtag = ""
 
         if torch.cuda.is_available():
             print('Found GPU will run ProteinMPNN on GPU')
@@ -84,6 +92,7 @@ def __init__(self, args, struct_manager):
         self.temperature = args.temperature
         self.seqs_per_struct = args.seqs_per_struct
         self.omit_AAs = [ letter for letter in args.omit_AAs.upper() if letter in list("ARNDCQEGHILKMFPSTWYVX") ]
+        print(f"The following amino acids will be omitted:  {self.omit_AAs}")
 
     def sequence_optimize(self, sample_feats: SampleFeatures) -> list[tuple[str, float]]:
         t0 = time.time()
@@ -128,13 +137,15 @@ def proteinmpnn(self, sample_feats: SampleFeatures) -> None:
 
         # Iterate though each seq score pair and thread the sequence onto the pose
         # Then write each pose to a pdb file
-        prefix = f"{sample_feats.tag}_dldesign"
+        prefix = f"{sample_feats.tag}_{args.name_tag}"
+        print("prefix",prefix)
+        print(seqs_scores)
         for idx, (seq, _) in enumerate(seqs_scores): 
             sample_feats.thread_mpnn_seq(seq)
 
-            outtag = f"{prefix}_{idx}"
+            self.outtag = f"{prefix}_{idx}"
 
-            self.struct_manager.dump_pose(sample_feats.pose, outtag)
+            self.struct_manager.dump_pose(sample_feats.pose, self.outtag)
 
     def run_model(self, tag, args):
         t0 = time.time()
@@ -161,6 +172,17 @@ def run_model(self, tag, args):
 ####### Main #######
 ####################
 
+if args.seed:
+    seed=int(args.seed)
+    if seed < 0 or seed > 999:
+        print("Seed has to be in the 0-999 range, 111 will be used!") 
+    print(f"Seed: {seed} will be used")
+else:
+    seed=int(np.random.randint(0, high=999, size=1, dtype=int)[0])
+    print("Random seed wil be used")
+torch.manual_seed(seed)
+
+
 struct_manager = StructManager(args)
 proteinmpnn_runner = ProteinMPNN_runner(args, struct_manager)
 
@@ -180,7 +202,7 @@ def run_model(self, tag, args):
             print(f"Struct with tag {pdb} failed in {seconds} seconds with error: {sys.exc_info()[0]}")
 
     # We are done with one pdb, record that we finished
-    struct_manager.record_checkpoint(pdb)
+    struct_manager.record_checkpoint(f"{args.outpdbdir}/{proteinmpnn_runner.outtag}.pdb")
 
 
 

scripts/rf2_predict.py

@@ -8,7 +8,7 @@
 from rfantibody.rf2.modules.model_runner import AbPredictor
 from rfantibody.rf2.modules.preprocess import pose_to_inference_RFinput, Preprocess
 
-@hydra.main(version_base=None, config_path='/home/src/rfantibody/rf2/config', config_name='base')
+@hydra.main(version_base=None, config_path='/opt/RFantibody/src/rfantibody/rf2/config', config_name='base')
 def main(conf: HydraConfig) -> None:
     """
     Main function

scripts/rfdiffusion_inference.py

@@ -42,7 +42,7 @@ def make_deterministic(seed=0):
         np.random.seed(seed)
         random.seed(seed)
 
-@hydra.main(version_base=None, config_path='config/inference', config_name='base')
+@hydra.main(version_base=None, config_path='/opt/RFantibody/src/rfantibody/rfdiffusion/config/inference', config_name='base')
 def main(conf: HydraConfig) -> None:
     log = logging.getLogger(__name__)
     if conf.inference.deterministic:
@@ -266,7 +266,8 @@ def main(conf: HydraConfig) -> None:
                 )
 
                 with open(out, 'w') as f_out:
-                    f_out.write('\n'.join(pdblines))
+                    # f_out.write('\n'.join(pdblines))  # this produced an empty lin after each line
+                    f_out.write(''.join(pdblines))
 
             else:
                 # Add to Quiver file