added threading and Colabfold-like plots generation

openfold/config.py

@@ -224,7 +224,7 @@ def model_config(
             c.data.eval.max_extra_msa = 1152
             c.data.predict.max_extra_msa = 1152
     else:
-        raise ValueError("Invalid model name")
+        raise ValueError(f"Invalid model name {name}")
 
     if long_sequence_inference:
         assert(not train)

openfold/data/data_pipeline.py

@@ -1208,7 +1208,7 @@ def read_msa(start, size):
             uniprot_msa_path = os.path.join(alignment_dir, "uniprot_hits.sto")
             if not os.path.exists(uniprot_msa_path):
                 chain_id = os.path.basename(os.path.normpath(alignment_dir))
-                raise ValueError(f"Missing 'uniprot_hits.sto' for {chain_id}. "
+                raise ValueError(f"Missing file {uniprot_msa_path} for {chain_id}. "
                                  f"This is required for Multimer MSA pairing.")
 
             with open(uniprot_msa_path, "r") as fp:
@@ -1235,7 +1235,6 @@ def process_fasta(self,
             input_fasta_str = f.read()
 
         input_seqs, input_descs = parsers.parse_fasta(input_fasta_str)
-
         all_chain_features = {}
         sequence_features = {}
         is_homomer_or_monomer = len(set(input_seqs)) == 1

openfold/data/tools/hhsearch.py

@@ -107,8 +107,8 @@ def query(self, a3m: str, output_dir: Optional[str] = None) -> str:
             if retcode:
                 # Stderr is truncated to prevent proto size errors in Beam.
                 raise RuntimeError(
-                    "HHSearch failed:\nstdout:\n%s\n\nstderr:\n%s\n"
-                    % (stdout.decode("utf-8"), stderr[:100_000].decode("utf-8"))
+                    "HHSearch failed:\ncommand:\n%s\n\nstdout:\n%s\n\nstderr:\n%s\n"
+                    % (f"hhsearch command: {' '.join(cmd)}", stdout.decode("utf-8"), stderr[:100_000].decode("utf-8"))
                 )
 
             with open(hhr_path) as f:

run_pretrained_openfold.py

@@ -138,7 +138,8 @@ def generate_feature_dict(
                 '\n'.join([f">{tag}\n{seq}" for tag, seq in zip(tags, seqs)])
             )
         feature_dict = data_processor.process_fasta(
-            fasta_path=tmp_fasta_path, alignment_dir=alignment_dir,
+            fasta_path=tmp_fasta_path,
+            alignment_dir=alignment_dir
         )
     elif len(seqs) == 1:
         tag = tags[0]
@@ -180,6 +181,45 @@ def main(args):
 
     config = model_config(args.config_preset, long_sequence_inference=args.long_sequence_inference)
 
+    print("")
+    print("#### INPUT / OUTPUT ####")
+    print(f"fasta_dir: {args.fasta_dir}")
+    print(f"output_dir: {args.output_dir}")
+    print(f"output prediction filenames: {args.output_postfix}")
+    print(f"cif_output: {args.cif_output}")
+    print(f"save embedded outputs: {args.save_outputs}")
+
+    print("")
+    print("#### PRESETS ####")
+    print(f"skip_relaxation: {args.skip_relaxation}")
+    print(f"use_precomputed_alignments: {args.use_precomputed_alignments}")
+    print(f"use_single_seq_mode: {args.use_single_seq_mode}")
+    print(f"long_sequence_inference: {args.long_sequence_inference}")
+    print(f"Threads: {args.cpus}")
+    print(f"multimer_ri_gap: {args.multimer_ri_gap}")
+    print(f"subtract_plddt: {args.subtract_plddt}")
+
+    print("")
+    print("#### MODEL PARAMS ####")
+    print(f"Model: {args.config_preset}")
+    print(f"trace_model: {args.trace_model}")
+
+    print("")
+    print("#### DATABASE PARAMS ####")
+    print(f"template_mmcif_dir: {args.template_mmcif_dir}")
+    print(f"max_template_date: {args.max_template_date}")
+    print(f"max_templates: {config.data.predict.max_templates}")
+    print(f"release_dates_path: {args.release_dates_path}")
+    print(f"obsolete_pdbs_path: {args.obsolete_pdbs_path}")
+
+    print("")
+    print("#### GPU / AI PARAMS ####")
+    print(f"model_device: {args.model_device}")
+    print(f"openfold_checkpoint_path: {args.openfold_checkpoint_path}")
+    print(f"jax_param_path: {args.jax_param_path}")
+
+    print("")
+
     if args.trace_model:
         if not config.data.predict.fixed_size:
             raise ValueError(
@@ -237,6 +277,7 @@ def main(args):
     for fasta_file in list_files_with_extensions(args.fasta_dir, (".fasta", ".fa")):
         # Gather input sequences
         fasta_path = os.path.join(args.fasta_dir, fasta_file)
+        print(f"reading fasta: {fasta_path}")
         with open(fasta_path, "r") as fp:
             data = fp.read()
 
@@ -258,12 +299,15 @@ def main(args):
     seq_sort_fn = lambda target: sum([len(s) for s in target[1]])
     sorted_targets = sorted(zip(tag_list, seq_list), key=seq_sort_fn)
     feature_dicts = {}
+
+    logger.info(f"Loading model information...")
     model_generator = load_models_from_command_line(
         config,
         args.model_device,
         args.openfold_checkpoint_path,
         args.jax_param_path,
-        args.output_dir)
+        args.output_dir
+    )
 
     for model, output_directory in model_generator:
         cur_tracing_interval = 0
@@ -273,6 +317,7 @@ def main(args):
                 output_name = f'{output_name}_{args.output_postfix}'
 
             # Does nothing if the alignments have already been computed
+            logger.info(f"Perform alignment if not already done...")
             precompute_alignments(tags, seqs, alignment_dir, args)
 
             feature_dict = feature_dicts.get(tag, None)
@@ -298,6 +343,10 @@ def main(args):
                 feature_dict, mode='predict', is_multimer=is_multimer
             )
 
+            # print("Storing feature dict...")
+            # with open(os.path.join(args.output_dir, f"{output_name}_feature_dict.pickle"), "wb") as fp:
+            #     pickle.dump(processed_feature_dict, fp, protocol=pickle.HIGHEST_PROTOCOL)
+
             processed_feature_dict = {
                 k: torch.as_tensor(v, device=args.model_device)
                 for k, v in processed_feature_dict.items()
@@ -316,6 +365,10 @@ def main(args):
                     )
                     cur_tracing_interval = rounded_seqlen
 
+            print("Storing feature dict...")
+            with open(os.path.join(args.output_dir, f"{output_name}_feature_dict.pickle"), "wb") as fp:
+                pickle.dump(feature_dict, fp, protocol=pickle.HIGHEST_PROTOCOL)
+
             out = run_model(model, processed_feature_dict, tag, args.output_dir)
 
             # Toss out the recycling dimensions --- we don't need them anymore