Batch per-genome processes

bin/generate_auxiliary_db.py

@@ -99,66 +99,67 @@ def sample_feature_to_annotation_entry(bakta_feature: dict) -> dict:
     return protein_annotation
 
 
-def update_cds_annotation(pickle_paths: Iterable[Path], cds_annotation_before_filtering: dict) -> dict:
-    for pickle_path in pickle_paths:
-        sample_data = ut.load_pickle(pickle_path)  # TODO: change if other serialization formats will be added
-        for feature in sample_data["features"]:
-            if feature["type"] not in CACHED_FEATURE_TYPES:
-                continue
+def update_cds_annotation(batch_pickle_paths: Iterable[Path], cds_annotation_before_filtering: dict) -> dict:
+    for pickle_path in batch_pickle_paths:
+        batch_data = ut.load_pickle(pickle_path)  # TODO: change if other serialization formats will be added
+        for sample_id, sample_data in batch_data.items():
+            for feature in sample_data["features"]:
+                if feature["type"] not in CACHED_FEATURE_TYPES:
+                    continue
 
-            protein_anno = sample_feature_to_annotation_entry(feature)
+                protein_anno = sample_feature_to_annotation_entry(feature)
 
-            if feature["aa_hexdigest"] not in cds_annotation_before_filtering:
-                cds_annotation_before_filtering[feature["aa_hexdigest"]] = protein_anno
-            else:
-                cds_annotation_before_filtering[feature["aa_hexdigest"]] |= protein_anno
+                if feature["aa_hexdigest"] not in cds_annotation_before_filtering:
+                    cds_annotation_before_filtering[feature["aa_hexdigest"]] = protein_anno
+                else:
+                    cds_annotation_before_filtering[feature["aa_hexdigest"]] |= protein_anno
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
-        description="Extract protein annotations from a Bakta JSON and write an index keyed by aa_hexdigest"
+        description="Extract protein annotations from serialised Bakta output data and write an index keyed by aa_hexdigest to generate a Pannotator auxiliary database"
     )
 
     parser.add_argument(
         "-c",
-        "--cds_annotation_pickles",
+        "--batch_cds_annotation_pickles",
         type=Path,
         nargs="+",
         required=True,
-        help="List of CDS annotation pickles with annotated pseudogenes",
+        help="List of batch CDS annotation pickles including pseudogenes",
     )
 
     parser.add_argument(
         "-i",
-        "--sorf_annotation_pickles",
+        "--batch_sorf_annotation_pickles",
         type=Path,
         nargs="+",
         required=True,
-        help="List of annotation pickle objects to extract sORF features from and generate auxiliary DB",
+        help="List of batch sORF annotation pickles",
     )
 
     parser.add_argument(
         "-a",
         "--auxiliary_db",
         type=Path,
         required=True,
-        help="Path to JSON to generate or update with new proteins",
+        help="Path to Pannotator auxiliary database (JSON format) to be generated or updated with new protein entries",
     )
 
     parser.add_argument(
         "-b",
         "--bulk_annotation_before_filtering",
         type=Path,
         required=True,
-        help="Path to bulk CDS annotation JSON",
+        help="Path to bulk CDS annotation (JSON format)",
     )
 
     parser.add_argument(
         "-o",
         "--updated_db_out",
         type=Path,
         required=True,
-        help="Output path for updated auxiliary JSON DB",
+        help="Output path for updated auxiliary database (JSON format)",
     )
 
     args = parser.parse_args()

bin/manage_pkls.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+"""Batch and unbatch Bakta pickle outputs (CDS/RNA).
+
+batch:   merge N per-sample pickles into one batch-level pickle, {sample_id: data, ...}
+unbatch: split one batch-level pickle back into N per-sample pickle files,
+         plus a manifest listing sample_id -> written path (for tools like
+         bakta_pseudo_bulk that expect a manifest of individual files).
+"""
+
+import argparse
+from pathlib import Path
+
+import utils as ut
+
+
+def batch_pickles(sample_ids: list[str], pkl_files: list[Path]) -> dict:
+    if len(sample_ids) != len(pkl_files):
+        raise ValueError(
+            f"Mismatch: {len(sample_ids)} sample IDs vs {len(pkl_files)} pickle files"
+        )
+    return {sid: ut.load_pickle(path) for sid, path in zip(sample_ids, pkl_files)}
+
+
+def unbatch_pickles(batch: dict, out_dir: Path, suffix: str) -> dict[str, Path]:
+    out_dir.mkdir(parents=True, exist_ok=True)
+    written = {}
+    for sample_id, data in batch.items():
+        out_path = out_dir / f"{sample_id}{suffix}"
+        ut.dump_pickle(data, out_path)
+        written[sample_id] = out_path
+    return written
+
+
+def write_manifest(paths_by_id: dict[str, Path], manifest_path: Path) -> None:
+    with open(manifest_path, "w") as mf:
+        mf.write("sample_id\tpath\n")
+        mf.writelines(f"{sample_id}\t{path}\n" for sample_id, path in paths_by_id.items())
+
+
+def cmd_batch(args):
+    sample_ids = args.sample_ids.split(",")
+    merged = batch_pickles(sample_ids, args.pkl_files)
+    ut.dump_pickle(merged, Path(args.output))
+    print(f"Batched {len(merged)} samples into {args.output}")
+
+
+def cmd_unbatch(args):
+    batch = ut.load_pickle(Path(args.input))
+    written = unbatch_pickles(batch, Path(args.out_dir), args.suffix)
+    print(f"Unbatched {len(written)} samples into {args.out_dir}")
+    if args.manifest:
+        write_manifest(written, Path(args.manifest))
+        print(f"Manifest written to {args.manifest}")
+
+
+def main():
+    p = argparse.ArgumentParser(description="Batch/unbatch Bakta pickle files")
+    sub = p.add_subparsers(dest="command", required=True)
+
+    batch_p = sub.add_parser("batch", help="Merge per-sample pickles into one batch pickle")
+    batch_p.add_argument("--sample-ids", required=True, help="Comma-separated sample IDs, positionally identical to pkl_files")
+    batch_p.add_argument("--output", required=True, help="Output batch pickle path")
+    batch_p.add_argument("pkl_files", nargs="+", type=Path, help="Per-sample pickle files, in sample-id order")
+    batch_p.set_defaults(func=cmd_batch)
+
+    unbatch_p = sub.add_parser("unbatch", help="Split a batch pickle back into per-sample files")
+    unbatch_p.add_argument("--input", required=True, help="Batch-level pickle file to split")
+    unbatch_p.add_argument("--out-dir", required=True, help="Directory to write per-sample pickle files")
+    unbatch_p.add_argument("--suffix", default=".pkl", help="Filename suffix for written files, e.g. '.cds-only.pkl'")
+    unbatch_p.add_argument("--manifest", default=None, help="Optional manifest TSV path (sample_id -> written path)")
+    unbatch_p.set_defaults(func=cmd_unbatch)
+
+    args = p.parse_args()
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

bin/merge_annotations_into_pkl.py

@@ -71,81 +71,86 @@ def find_annotation(feature: Dict[str, Any], annotations: Dict[str, Dict[str, An
 
     return True
 
-
-def process_sample(pkl_path: Path, annotations: Dict[str, Dict[str, Any]], output_path: Path) -> None:
-    print(f"Processing sample: {pkl_path}")
-
-    total = 0
-    annotated = 0
-    missing = []
-
-    data = ut.load_pickle(pkl_path)  # TODO: change if more serialization formats will be added
-
+def annotate_sample_features(sample_id: str, data: Dict[str, Any], annotations: Dict[str, Dict[str, Any]]) -> None:
     features = data.get("features", [])
     if not features:
-        print(f"No features found in {pkl_path}")
+        print(f"No features found for sample {sample_id}")
         return
 
+    total = annotated = 0
+    missing = []
     for feature in features:
         if feature.get("type") == "cds":
             total += 1
-
             if find_annotation(feature, annotations):
                 annotated += 1
             else:
                 missing.append(feature.get("aa_hexdigest", "no_hexdigest"))
 
-    print(f"Total CDS features: {total}")
-    print(f"Annotated: {annotated}")
-    print(f"Missing: {len(missing)}")
-
-    ut.dump_pickle(data, output_path)  # TODO: change if more serialization formats will be added
+    print(f"[{sample_id}] CDS features: {total}, annotated: {annotated}, missing: {len(missing)}")
 
-    print(f"Annotated pickle saved to: {output_path}")
+def process_batch(pkl_path: Path, annotations: dict, output_path: Path) -> None:
+    print(f"Processing batch: {pkl_path}")
+    batch = ut.load_pickle(pkl_path)  # {sample_id: {"features": [...], ...}, ...}
 
+    if not batch:
+        print(f"No samples found in {pkl_path}")
+        return
 
-def process_input_folder(input_folder: Path, json_path: Path, output_folder: Path) -> None:
-    annotations = load_annotations(json_path)
+    for sample_id, data in batch.items():
+        annotate_sample_features(sample_id, data, annotations)  # unchanged, mutates in place
 
-    pickle_files = list(input_folder.glob("*.cds-only.pkl"))  # TODO: change if more serialization formats will be added
+    ut.dump_pickle(batch, output_path)
+    print(f"Annotated batch pickle saved to: {output_path}")
 
-    if not pickle_files:
-        print(f"No sample pickle files found in {input_folder}")
-        return
 
-    print(f"Found {len(pickle_files)} sample pickle files to propagate annotations to.")
+def process_input_pickle(input_pickle: Path, json_path: Path, output_folder: Path) -> None:
+    annotations = load_annotations(json_path)
+    new_name = input_pickle.name.replace(".cds-only.pkl", ".cds-annotated.pkl")
+    output_path = output_folder / new_name
+    process_batch(input_pickle, annotations, output_path)
 
-    output_folder.mkdir(parents=True, exist_ok=True)
+# def process_input_folder(input_folder: Path, json_path: Path, output_folder: Path) -> None:
+#     annotations = load_annotations(json_path)
+#     pickle_files = list(input_folder.glob("*.cds-only.pkl"))
 
-    for pkl_path in pickle_files:
-        new_name = pkl_path.name.replace(".cds-only.pkl", ".cds-annotated.pkl")
-        output_path = output_folder / new_name
+#     if not pickle_files:
+#         print(f"No batch pickle files found in {input_folder}")
+#         return
 
-        process_sample(pkl_path, annotations, output_path)
+#     print(f"Found {len(pickle_files)} batch pickle files to propagate annotations to.")
+#     output_folder.mkdir(parents=True, exist_ok=True)
 
-    print(f"All {len(pickle_files)} samples processed.")
+#     for pkl_path in pickle_files:
+#         new_name = pkl_path.name.replace(".cds-only.pkl", ".cds-annotated.pkl")
+#         output_path = output_folder / new_name
+#         process_batch(pkl_path, annotations, output_path)
 
+#     print(f"All {len(pickle_files)} batches processed.")
 
 def main():
-    p = argparse.ArgumentParser(description="Annotate CDS pickles with information from JSON annotations")
-    p.add_argument("--pickle_folder", required=True, help="Folder containing .cds-only.pkl files")
+    p = argparse.ArgumentParser(description="Annotate CDS pickles (batch format) with information from JSON annotations")
+    p.add_argument("--pickle_in", required=True, help="batch-level .cds-only.pkl file")
     p.add_argument("--annotations", required=True, help="bulk_protein_annotations.json file")
-    p.add_argument("--out", default="annotated_pkl", help="Output directory")
+    p.add_argument("--pickle_out", default="annotated_pkl", help="Folder for output annotated batch pickle files")
     args = p.parse_args()
 
-    input_folder = Path(args.pickle_folder)
+    input_pickle = Path(args.pickle_in)
     json_path = Path(args.annotations)
-    output_folder = Path(args.out)
+    output_folder = Path(args.pickle_out)
 
-    if not input_folder.exists():
-        print(f"Input folder does not exist: {input_folder}")
+    if not input_pickle.exists():
+        print(f"Input file does not exist: {input_pickle}")
         return
-
     if not json_path.exists():
         print(f"Annotations JSON file does not exist: {json_path}")
         return
 
-    process_input_folder(input_folder, json_path, output_folder)
+    if not output_folder.exists():
+        output_folder.mkdir(parents=True, exist_ok=True)
+
+    process_input_pickle(input_pickle, json_path, output_folder)
+
 
 
 if __name__ == "__main__":

bin/merge_pickles.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+"""Merge Bakta output pickle files for a batch of genomes of any size into a single batch-level pickle file"""
+
+import argparse
+from pathlib import Path
+
+import utils as ut
+
+
+def merge_batch(assembly_ids: list[str], pkl_files: list[Path]) -> dict:
+    if len(assembly_ids) != len(pkl_files):
+        raise ValueError(
+            f"Mismatch: {len(assembly_ids)} sample IDs vs {len(pkl_files)} pickle files"
+        )
+
+    merged = {}
+    for assembly_id, pkl_path in zip(assembly_ids, pkl_files):
+        merged[assembly_id] = ut.load_pickle(pkl_path)
+
+    return merged
+
+def main():
+    p = argparse.ArgumentParser(description="Merge pickle files within batches.")
+
+    p.add_argument(
+        "--assembly_ids", 
+        required=True, 
+        help="IDs for assemblies within the batch, comma-separated and positionally identical to pkls"
+        )
+    p.add_argument(
+        "--out",
+        required=True,
+        help="Output batch pickle file path"
+        )
+    p.add_argument(
+        "pkl_files",    # positional arg
+        nargs='+',
+        help="Pickle file paths, positionally identical to supplied IDs"
+        )
+
+    args = p.parse_args()
+
+    asm_ids = args.assembly_ids.split(',')
+    merged = merge_batch(asm_ids, args.pkl_files)
+    ut.dump_pickle(merged, Path(args.out))
+    print(f"Merged {len(merged)} samples into {args.out}")
+
+
+if __name__ == "__main__":
+    main()