| """Prepare the pinned BGC clustering benchmark without altering upstream files.""" |
|
|
| from __future__ import annotations |
|
|
| import ast |
| from pathlib import Path |
| from typing import Any |
|
|
| import pandas as pd |
| from Bio import SeqIO |
| from Bio.Seq import Seq |
| from Bio.SeqRecord import SeqRecord |
|
|
| from .artifacts import file_record, sha256_file, write_json_immutable |
|
|
|
|
| def build_product_mappings(source_tsv: str | Path) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: |
| source = pd.read_csv(source_tsv, sep="\t") |
| required = {"npaid", "compound_inchikey", "genus", "mibig_ids"} |
| if missing := required.difference(source.columns): |
| raise ValueError(f"NPAtlas source is missing columns: {sorted(missing)}") |
| rows: list[dict[str, str]] = [] |
| for record in source.itertuples(index=False): |
| try: |
| mibig_ids = ast.literal_eval(record.mibig_ids) if isinstance(record.mibig_ids, str) else [] |
| except (SyntaxError, ValueError) as error: |
| raise ValueError(f"Invalid mibig_ids for {record.npaid}") from error |
| for bgc_id in mibig_ids: |
| rows.append( |
| { |
| "bgc_id": str(bgc_id), |
| "product_id": str(record.compound_inchikey), |
| "genus": str(record.genus), |
| "npaid": str(record.npaid), |
| } |
| ) |
| expanded = pd.DataFrame(rows).drop_duplicates() |
| metadata_rows: list[dict[str, Any]] = [] |
| gold_rows: list[dict[str, str]] = [] |
| exclusions: list[dict[str, str]] = [] |
| for bgc_id, group in expanded.groupby("bgc_id"): |
| products = sorted(group["product_id"].dropna().unique()) |
| genera = sorted(group["genus"].dropna().unique()) |
| metadata_rows.append( |
| { |
| "bgc_id": bgc_id, |
| "product_ids": ";".join(products), |
| "product_count": len(products), |
| "genera": ";".join(genera), |
| "genus_count": len(genera), |
| } |
| ) |
| if len(products) == 1 and len(genera) == 1: |
| gold_rows.append( |
| { |
| "bgc_id": bgc_id, |
| "product_group_id": f"inchikey:{products[0]}", |
| "product_id": products[0], |
| "mibig_reference_id": bgc_id, |
| "genus": genera[0], |
| "source": "BGC-clustering-benchmark@bb8500d60f90", |
| } |
| ) |
| else: |
| reason = "multiple_products" if len(products) != 1 else "multiple_or_missing_genera" |
| exclusions.append( |
| {"bgc_id": bgc_id, "reason": reason, "values": ";".join(products)} |
| ) |
| return pd.DataFrame(metadata_rows), pd.DataFrame(gold_rows), pd.DataFrame(exclusions) |
|
|
|
|
| def _feature_translation(feature: Any, nucleotide_record: Any) -> str: |
| translations = feature.qualifiers.get("translation", []) |
| if translations: |
| return str(translations[0]).replace(" ", "").rstrip("*") |
| table = int(feature.qualifiers.get("transl_table", [11])[0]) |
| nucleotide = feature.extract(nucleotide_record.seq) |
| return str(nucleotide.translate(table=table, to_stop=False)).rstrip("*") |
|
|
|
|
| def extract_benchmark_proteins( |
| genbank_dir: str | Path, |
| fasta_path: str | Path, |
| atlas_path: str | Path, |
| provenance_path: str | Path, |
| source_commit: str, |
| ) -> dict[str, Any]: |
| fasta_output = Path(fasta_path) |
| atlas_output = Path(atlas_path) |
| provenance_output = Path(provenance_path) |
| for output in (fasta_output, atlas_output, provenance_output): |
| if output.exists(): |
| raise FileExistsError(f"Refusing to overwrite external preparation output: {output}") |
| output.parent.mkdir(parents=True, exist_ok=True) |
|
|
| protein_records: list[SeqRecord] = [] |
| atlas_rows: list[dict[str, Any]] = [] |
| excluded_files: list[dict[str, str]] = [] |
| source_files = sorted(Path(genbank_dir).glob("*.gbk")) |
| for genbank_path in source_files: |
| bgc_id = genbank_path.stem |
| protein_ids: list[str] = [] |
| organisms: set[str] = set() |
| ordinal = 0 |
| try: |
| with genbank_path.open("r", encoding="utf-8") as handle: |
| nucleotide_records = list(SeqIO.parse(handle, "genbank")) |
| for nucleotide_record in nucleotide_records: |
| organism = str(nucleotide_record.annotations.get("organism", "unknown")) |
| organisms.add(organism) |
| for feature in nucleotide_record.features: |
| if feature.type != "CDS": |
| continue |
| sequence = _feature_translation(feature, nucleotide_record) |
| if not sequence: |
| continue |
| gene_id = f"{bgc_id}__cds{ordinal:05d}" |
| protein_ids.append(gene_id) |
| protein_records.append( |
| SeqRecord(Seq(sequence), id=gene_id, description=f"source={bgc_id}") |
| ) |
| ordinal += 1 |
| except Exception as error: |
| excluded_files.append({"bgc_id": bgc_id, "reason": f"parse_error:{type(error).__name__}"}) |
| continue |
| if not protein_ids: |
| excluded_files.append({"bgc_id": bgc_id, "reason": "no_translated_cds"}) |
| continue |
| atlas_rows.append( |
| { |
| "bgc_id": bgc_id, |
| "protein_ids": ";".join(protein_ids), |
| "num_proteins": len(protein_ids), |
| "organisms": ";".join(sorted(organisms)), |
| } |
| ) |
|
|
| SeqIO.write(protein_records, fasta_output, "fasta-2line") |
| pd.DataFrame(atlas_rows).to_csv(atlas_output, index=False) |
| provenance = { |
| "schema_version": 1, |
| "source_commit": source_commit, |
| "source_files": len(source_files), |
| "usable_bgcs": len(atlas_rows), |
| "proteins": len(protein_records), |
| "excluded_files": excluded_files, |
| "fasta_sha256": sha256_file(fasta_output), |
| "atlas_sha256": sha256_file(atlas_output), |
| "source_manifest": [file_record(path, genbank_dir) for path in source_files], |
| } |
| write_json_immutable(provenance_output, provenance) |
| return provenance |
|
|
|
|
| def write_mapping_outputs( |
| source_tsv: str | Path, |
| mapping_path: str | Path, |
| metadata_path: str | Path, |
| exclusions_path: str | Path, |
| ) -> dict[str, int]: |
| metadata, mapping, exclusions = build_product_mappings(source_tsv) |
| for path in (mapping_path, metadata_path, exclusions_path): |
| output = Path(path) |
| output.parent.mkdir(parents=True, exist_ok=True) |
| if output.exists(): |
| raise FileExistsError(f"Refusing to overwrite external mapping output: {output}") |
| mapping.to_csv(mapping_path, index=False) |
| metadata.to_csv(metadata_path, index=False) |
| exclusions.to_csv(exclusions_path, index=False) |
| return { |
| "all_mapped_bgcs": len(metadata), |
| "unambiguous_gold_bgcs": len(mapping), |
| "ambiguous_exclusions": len(exclusions), |
| } |
|
|
