source/scripts/prepare_external_benchmark.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path

from bgc_retrieval.external_prepare import extract_benchmark_proteins, write_mapping_outputs


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--benchmark-dir", default="data/external/bgc-clustering-benchmark"
    )
    parser.add_argument("--output-dir", default="data/external/processed")
    parser.add_argument("--source-commit", default="bb8500d60f90cb43397cc41de5aed396725aa800")
    args = parser.parse_args()
    benchmark = Path(args.benchmark_dir)
    output = Path(args.output_dir)
    mapping_counts = write_mapping_outputs(
        benchmark / "source_data/NPAtlas_bm_v1.tsv",
        output / "gold_bgc_product_mapping.csv",
        output / "all_bgc_product_metadata.csv",
        output / "ambiguous_product_exclusions.csv",
    )
    sequence_counts = extract_benchmark_proteins(
        benchmark / "benchmark_bgc",
        output / "external_proteins.fasta",
        output / "external_atlas.csv",
        output / "external_sequence_provenance.json",
        args.source_commit,
    )
    concise_sequences = {
        key: value for key, value in sequence_counts.items() if key != "source_manifest"
    }
    print(json.dumps({**mapping_counts, **concise_sequences}, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()