Fixing model generation

run_pretrained_openfold.py

@@ -19,6 +19,7 @@ import gc
 import logging
 import numpy as np
 import os
+from copy import deepcopy
 
 import pickle
 from pytorch_lightning.utilities.deepspeed import (
@@ -289,14 +290,15 @@ def main(args):
 
         for model, model_version in load_models_from_command_line(args, config):
 
-            out = run_model(model, batch, tag, args)
+            working_batch = deepcopy(batch)
+            out = run_model(model, working_batch, tag, args)
 
             # Toss out the recycling dimensions --- we don't need them anymore
-            batch = tensor_tree_map(lambda x: np.array(x[..., -1].cpu()), batch)
+            working_batch = tensor_tree_map(lambda x: np.array(x[..., -1].cpu()), working_batch)
             out = tensor_tree_map(lambda x: np.array(x.cpu()), out)
 
             unrelaxed_protein = prep_output(
-                out, batch, feature_dict, feature_processor, args
+                out, working_batch, feature_dict, feature_processor, args
             )
 
             output_name = f'{tag}_{args.model_name}'

test.sh

-#!/bin/bash
-
-# the alphafold script has a few frustrating behaviors in how it parses input flags, this script reduces the need for
-# code duplication in the argo workflow
-set -x
-
-database_path=$1
-fasta_path=$2
-
-base_arguments=(
-    ${fasta_path}
-    "${database_path}/pdb_mmcif/mmcif_files/"
-    "--uniref90_database_path" "${database_path}/uniref90/uniref90.fasta"
-    "--mgnify_database_path" "${database_path}/mgnify/mgy_clusters_2018_12.fa"
-    "--pdb70_database_path" "${database_path}/pdb70/pdb70"
-    "--uniclust30_database_path" "${database_path}/uniclust30/uniclust30_2018_08/uniclust30_2018_08"
-    "--output_dir" "out"
-    "--bfd_database_path" "${database_path}/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt"
-    "--model_device" "cuda:0"
-    "--jackhmmer_binary_path" "/home/samdeluca/miniconda3/envs/openfold_venv/bin/jackhmmer"
-    "--hhblits_binary_path" "/home/samdeluca/miniconda3/envs/openfold_venv/bin/hhblits"
-    "--hhsearch_binary_path" "/home/samdeluca/miniconda3/envs/openfold_venv/bin/hhsearch"
-    "--kalign_binary_path" "/home/samdeluca/miniconda3/envs/openfold_venv/bin/kalign"
-    "--openfold_checkpoint_path" "/mnt/openfold_params/101-80999.ckpt,/mnt/openfold_params/116-84749.ckpt,/mnt/openfold_params/94-79249.ckpt"
-)
-
-python3 run_pretrained_openfold.py  ${base_arguments[*]}