alphafold_non_docker

--- /apps/alphafold/run_alphafold.sh.orig       2022-04-04 17:49:07.215185888 +0900
+++ /apps/alphafold/run_alphafold.sh    2022-04-04 17:48:57.735109552 +0900
@@ -131,7 +131,7 @@
 fi
 
 # This bash script looks for the run_alphafold.py script in its current working directory, if it does not exist then exits
-current_working_dir=$(pwd)
+current_working_dir=$alphafold_path
 alphafold_script="$current_working_dir/run_alphafold.py"
 
 if [ ! -f "$alphafold_script" ]; then

--- ../src/alphafold_non_docker/run_alphafold.sh.orig   2022-06-09 02:34:27.897005704 +0900
+++ run_alphafold.sh    2022-06-12 14:46:32.518462539 +0900
@@ -23,10 +23,15 @@
         echo "-l <num_multimer_predictions_per_model> How many predictions (each with a different random seed) will be (略
         echo "-b <benchmark>        Run multiple JAX model evaluations to obtain a timing that excludes the compilation (略
         echo ""
+        echo "-C <num_recycle>      ReCycle number [3]"
+        echo "-N <n_cpu>            jackhmmer: number of parallel CPU workers to use for multithreads [8]"
+        echo "-h <hhsearch_cpu>     hhsearch: number of CPUs to use (for shared memory SMPs) [2](monomer)"
+        echo "-H <hmmsearch_cpu>    hmmsearch: number of parallel CPU workers to use for multithreads [8](multimer)"
+        echo ""
         exit 1
 }
 
-while getopts ":d:o:f:t:g:r:e:n:a:m:c:p:l:b" i; do
+while getopts ":d:o:f:t:g:r:e:n:a:m:c:p:l:C:N:h:H:b" i; do
         case "${i}" in
         d)
                 data_dir=$OPTARG
@@ -67,6 +72,18 @@
         l)
                 num_multimer_predictions_per_model=$OPTARG
         ;;
+        C)
+                num_recycle=$OPTARG
+        ;;
+        N)
+                n_cpu=$OPTARG
+        ;;
+        h)
+                hhsearch_cpu=$OPTARG
+        ;;
+        H)
+                hmmsearch_cpu=$OPTARG
+        ;;
         b)
                 benchmark=true
         ;;
@@ -78,6 +95,18 @@
     usage
 fi
 
+if [[ "$hmmsearch_cpu" == "" ]] ; then
+    hmmsearch_cpu=8
+fi
+if [[ "$hhsearch_cpu" == "" ]] ; then
+    hhsearch_cpu=2
+fi
+if [[ "$n_cpu" == "" ]] ; then
+    n_cpu=8
+fi
+if [[ "$num_recycle" == "" ]] ; then
+    num_recycle=3
+fi
 if [[ "$benchmark" == "" ]] ; then
     benchmark=false
 fi
@@ -131,7 +160,7 @@
 fi
 
 # This bash script looks for the run_alphafold.py script in its current working directory, if it does not exist then exits
-current_working_dir=$(pwd)
+current_working_dir=$alphafold_path
 alphafold_script="$current_working_dir/run_alphafold.py"
 
 if [ ! -f "$alphafold_script" ]; then
@@ -197,5 +226,6 @@
        database_paths="$database_paths --uniclust30_database_path=$uniclust30_database_path --bfd_database_path=$bfd_database_path"
 fi
 
+extra_args="--num_recycle=$num_recycle --n_cpu=$n_cpu --hhsearch_cpu=$hhsearch_cpu --hmmsearch_cpu=$hmmsearch_cpu"
 # Run AlphaFold with required parameters
-$(python $alphafold_script $binary_paths $database_paths $command_args)
+$(python $alphafold_script $binary_paths $database_paths $command_args $extra_args)

--- run_alphafold.py.orig       2022-06-09 02:35:22.146479521 +0900
+++ run_alphafold.py    2022-06-12 14:43:24.842855059 +0900
@@ -128,6 +128,10 @@
                      'Relax on GPU can be much faster than CPU, so it is '
                      'recommended to enable if possible. GPUs must be available'
                      ' if this setting is enabled.')
+flags.DEFINE_integer('num_recycle', None,'num_recycle')
+flags.DEFINE_integer('n_cpu', 8,'n_cpu')
+flags.DEFINE_integer('hhsearch_cpu', 2,'hhsearch_cpu')
+flags.DEFINE_integer('hmmsearch_cpu', 8,'hmmsearch_cpu')
 
 FLAGS = flags.FLAGS
 
@@ -315,6 +319,7 @@
     template_searcher = hmmsearch.Hmmsearch(
         binary_path=FLAGS.hmmsearch_binary_path,
         hmmbuild_binary_path=FLAGS.hmmbuild_binary_path,
+        hmmsearch_cpu=FLAGS.hmmsearch_cpu,
         database_path=FLAGS.pdb_seqres_database_path)
     template_featurizer = templates.HmmsearchHitFeaturizer(
         mmcif_dir=FLAGS.template_mmcif_dir,
@@ -326,6 +331,7 @@
   else:
     template_searcher = hhsearch.HHSearch(
         binary_path=FLAGS.hhsearch_binary_path,
+        hhsearch_cpu=FLAGS.hhsearch_cpu,
         databases=[FLAGS.pdb70_database_path])
     template_featurizer = templates.HhsearchHitFeaturizer(
         mmcif_dir=FLAGS.template_mmcif_dir,
@@ -337,6 +343,7 @@
 
   monomer_data_pipeline = pipeline.DataPipeline(
       jackhmmer_binary_path=FLAGS.jackhmmer_binary_path,
+      n_cpu=FLAGS.n_cpu,
       hhblits_binary_path=FLAGS.hhblits_binary_path,
       uniref90_database_path=FLAGS.uniref90_database_path,
       mgnify_database_path=FLAGS.mgnify_database_path,
@@ -359,6 +366,10 @@
     num_predictions_per_model = 1
     data_pipeline = monomer_data_pipeline
 
+  num_recycle = FLAGS.num_recycle
+  if num_recycle is None:
+    num_recycle = 3
+
   model_runners = {}
   model_names = config.MODEL_PRESETS[FLAGS.model_preset]
   for model_name in model_names:
@@ -367,6 +378,7 @@
       model_config.model.num_ensemble_eval = num_ensemble
     else:
       model_config.data.eval.num_ensemble = num_ensemble
+    model_config.data.common.num_recycle = FLAGS.num_recycle
     model_params = data.get_model_haiku_params(
         model_name=model_name, data_dir=FLAGS.data_dir)
     model_runner = model.RunModel(model_config, model_params)
@@ -417,6 +429,7 @@
       'max_template_date',
       'obsolete_pdbs_path',
       'use_gpu_relax',
+      'num_recycle',
   ])
 
   app.run(main)

--- a/alphafold/data/pipeline.py
+++ b/alphafold/data/pipeline.py
@@ -124,15 +124,18 @@ class DataPipeline:
                use_small_bfd: bool,
                mgnify_max_hits: int = 501,
                uniref_max_hits: int = 10000,
+               n_cpu: int = 8,
                use_precomputed_msas: bool = False):
     """Initializes the data pipeline."""
     self._use_small_bfd = use_small_bfd
     self.jackhmmer_uniref90_runner = jackhmmer.Jackhmmer(
         binary_path=jackhmmer_binary_path,
+        n_cpu=n_cpu,
         database_path=uniref90_database_path)
     if use_small_bfd:
       self.jackhmmer_small_bfd_runner = jackhmmer.Jackhmmer(
           binary_path=jackhmmer_binary_path,
+          n_cpu=n_cpu,
           database_path=small_bfd_database_path)
     else:
       self.hhblits_bfd_uniclust_runner = hhblits.HHBlits(
@@ -140,6 +143,7 @@ class DataPipeline:
           databases=[bfd_database_path, uniclust30_database_path])
     self.jackhmmer_mgnify_runner = jackhmmer.Jackhmmer(
         binary_path=jackhmmer_binary_path,
+        n_cpu=n_cpu,
         database_path=mgnify_database_path)
     self.template_searcher = template_searcher
     self.template_featurizer = template_featurizer

--- a/alphafold/data/tools/hhsearch.py
+++ b/alphafold/data/tools/hhsearch.py
@@ -33,6 +33,7 @@ class HHSearch:
                *,
                binary_path: str,
                databases: Sequence[str],
+               hhsearch_cpu: int = 2,
                maxseq: int = 1_000_000):
     """Initializes the Python HHsearch wrapper.
 
@@ -50,6 +51,7 @@ class HHSearch:
     self.binary_path = binary_path
     self.databases = databases
     self.maxseq = maxseq
+    self.hhsearch_cpu = hhsearch_cpu
 
     for database_path in self.databases:
       if not glob.glob(database_path + '_*'):
@@ -79,6 +81,7 @@ class HHSearch:
       cmd = [self.binary_path,
              '-i', input_path,
              '-o', hhr_path,
+             '-cpu', str(self.hhsearch_cpu),
              '-maxseq', str(self.maxseq)
              ] + db_cmd
 

--- a/alphafold/data/tools/hmmsearch.py
+++ b/alphafold/data/tools/hmmsearch.py
@@ -33,6 +33,7 @@ class Hmmsearch(object):
                binary_path: str,
                hmmbuild_binary_path: str,
                database_path: str,
+               hmmsearch_cpu: int = 8,
                flags: Optional[Sequence[str]] = None):
     """Initializes the Python hmmsearch wrapper.
 
@@ -49,6 +50,7 @@ class Hmmsearch(object):
     self.binary_path = binary_path
     self.hmmbuild_runner = hmmbuild.Hmmbuild(binary_path=hmmbuild_binary_path)
     self.database_path = database_path
+    self.hmmsearch_cpu = hmmsearch_cpu
     if flags is None:
       # Default hmmsearch run settings.
       flags = ['--F1', '0.1',
@@ -89,7 +91,7 @@ class Hmmsearch(object):
       cmd = [
           self.binary_path,
           '--noali',  # Don't include the alignment in stdout.
-          '--cpu', '8'
+          '--cpu', str(self.hmmsearch_cpu)
       ]
       # If adding flags, we have to do so before the output and input:
       if self.flags: