Add script for generating sample feature dict

scripts/generate_alphafold_feature_dict.py

+import argparse
+import os
+import pickle
+
+from alphafold.data import pipeline, templates
+
+from scripts.utils import add_data_args
+
+
+def main(args):
+    template_featurizer = templates.TemplateHitFeaturizer(
+        mmcif_dir=args.mmcif_dir,
+        max_template_date=args.max_template_date,
+        max_hits=20,
+        kalign_binary_path=args.kalign_binary_path,
+        release_dates_path=None,
+        obsolete_pdbs_path=args.obsolete_pdbs_path,
+    )
+
+    data_pipeline = pipeline.DataPipeline(
+        jackhmmer_binary_path=args.jackhmmer_binary_path,
+        hhblits_binary_path=args.hhblits_binary_path,
+        hhsearch_binary_path=args.hhsearch_binary_path,
+        uniref90_database_path=args.uniref90_database_path,
+        mgnify_database_path=args.mgnify_database_path,
+        bfd_database_path=args.bfd_database_path,
+        uniclust30_database_path=args.uniclust30_database_path,
+        pdb70_database_path=args.pdb70_database_path,
+        small_bfd_database_path=None,
+        template_featurizer=template_featurizer,
+        use_small_bfd=False,
+    )
+
+    feature_dict = data_pipeline.process(
+        input_fasta_path=args.fasta_path,
+        msa_output_dir=args.output_dir,
+    )
+
+    with open(os.path.join(args.output_dir, "feature_dict.pickle"), "wb") as fp:
+        pickle.dump(feature_dict, fp, protocol=pickle.HIGHEST_PROTOCOL)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("fasta_path", type=str)
+    parser.add_argument("mmcif_dir", type=str)
+    parser.add_argument("output_dir", type=str)
+    add_data_args(parser)
+
+    args = parser.parse_args()
+
+    main(args)