diff --git a/bin/count_phrog_functions.py b/bin/count_phrog_functions.py index f1c4212..3bbb498 100644 --- a/bin/count_phrog_functions.py +++ b/bin/count_phrog_functions.py @@ -1,19 +1,94 @@ """Summarise PHROG-annotated MMseqs top-hit reports across samples.""" import argparse +import csv import gzip import math import os +import re import sys from collections import defaultdict REPORT_SUFFIX = "_tophit_report_phrog_function.tsv.gz" +MARKER_ORDER = ( + "all_phrog_hits", + "known_function_phrog_hits", + "unknown_function_phrog_hits", + "integrase", + "excisionase", + "integrase_plus_excisionase", + "lysogeny_associated_regulators", + "major_capsid", + "portal", + "terminase", + "tail", + "holin", + "endolysin", + "holin_plus_endolysin", +) + +# Marker rules are intentionally conservative and centralized here for review. +EXCISIONASE_EXACT = { + "excisionase", + "cox-like excisionase and repressor", + "excisionase and transcriptional regulator", + "recombination directionality factor", + "site-specific recombination directionality factor rdf", +} +LYSOGENY_REGULATOR_EXACT = { + "ci-like repressor", + "cro-like repressor", + "cox-like excisionase and repressor", + "arc-like repressor", + "arc-like transcriptional regulator", + "copg-like transcriptional repressor", + "transcriptional repressor", + "anti-repressor", + "anti-repressor ant", + "cii-like regulator", + "cii-like transcriptional activator", + "ciii anti-termination", + "immunity protein", + "immunity to superinfection", + "excisionase and transcriptional regulator", +} +MAJOR_CAPSID_EXACT = { + "major capsid protein", + "major head protein", + "major head and protease protein", + "major coat protein", +} +ENDOLYSIN_EXACT = { + "endolysin", + "endolysin; inhibits rna polymerase", + "endolysin; lysm motif", + "lysozyme domain-containing protein", + "amidase", + "internal virion protein with endolysin domain", + "baseplate hub and tail lysozyme", + "baseplate hub subunit and tail lysozyme", + "tail associated lysin", + "tail fiber protein/ lysozyme", + "tail protein with lysin activity", + "minor tail protein with lysin activity", +} +UNKNOWN_ANNOTATIONS = { + "", + "-", + "na", + "n/a", + "none", + "unknown", + "unknown function", + "hypothetical protein", + "uncharacterized protein", +} def parse_args(): parser = argparse.ArgumentParser( - description="Create raw and counts-per-million PHROG summary tables" + description="Create raw and counts-per-million PHROG and marker tables" ) parser.add_argument( "-d", @@ -30,10 +105,127 @@ def parse_args(): parser.add_argument( "-n", "--name", default="atavide", help="prefix for output files [atavide]" ) + parser.add_argument( + "-a", + "--phrog-annotations", + required=True, + help=( + "PHROG annotation TSV with phrog, annot, and category columns; " + "plain text or gzip-compressed (for example phrog_annot_v4.tsv.gz)" + ), + ) parser.add_argument("-v", "--verbose", action="store_true") return parser.parse_args() +def normalise_text(value): + return " ".join(value.strip().split()).casefold() + + +def normalise_phrog_id(value): + value = value.strip() + match = re.fullmatch(r"(?:phrog_)?(\d+)", value, flags=re.IGNORECASE) + if not match: + raise ValueError( + f"invalid PHROG identifier {value!r}; expected a number or phrog_" + ) + return str(int(match.group(1))) + + +def open_text(path): + if path.endswith(".gz"): + return gzip.open(path, "rt", encoding="utf-8", newline="") + return open(path, "r", encoding="utf-8", newline="") + + +def read_phrog_annotations(path): + if not os.path.isfile(path): + raise SystemExit(f"FATAL: PHROG annotation database does not exist: {path}") + if os.path.getsize(path) == 0: + raise SystemExit(f"FATAL: PHROG annotation database is empty: {path}") + try: + with open_text(path) as input_file: + reader = csv.DictReader(input_file, delimiter="\t") + required = {"phrog", "annot", "category"} + if reader.fieldnames is None or not required.issubset(reader.fieldnames): + found = ", ".join(reader.fieldnames or []) or "no header" + raise SystemExit( + f"FATAL: {path} must contain tab-separated columns phrog, annot, " + f"and category (found: {found})" + ) + records = {} + for line_number, row in enumerate(reader, start=2): + try: + phrog_id = normalise_phrog_id(row["phrog"]) + except (AttributeError, ValueError) as error: + raise SystemExit(f"FATAL: {path}:{line_number}: {error}") from None + annotation = (row["annot"] or "").strip() + category = (row["category"] or "").strip() + record = (annotation, category) + if phrog_id in records and records[phrog_id] != record: + raise SystemExit( + f"FATAL: {path}:{line_number} has a conflicting record for " + f"PHROG {phrog_id}: {records[phrog_id]!r} versus {record!r}" + ) + records[phrog_id] = record + except (OSError, UnicodeError, csv.Error) as error: + raise SystemExit(f"FATAL: cannot read PHROG annotation database {path}: {error}") from None + if not records: + raise SystemExit(f"FATAL: PHROG annotation database contains no records: {path}") + return records + + +def classify_phrog(annotation, category): + """Return {marker: auditable rule} for one authoritative PHROG record.""" + annot = normalise_text(annotation) + cat = normalise_text(category) + markers = {} + + if re.search(r"\bintegrase\b", annot): + markers["integrase"] = "word-bounded annotation match: integrase" + if annot in EXCISIONASE_EXACT: + markers["excisionase"] = "curated exact excisionase/RDF annotation" + if annot in LYSOGENY_REGULATOR_EXACT or re.search( + r"\b(?:repressor|anti-repressor|antirepressor)\b", annot + ) or annot in {"immunity protein", "immunity to superinfection"}: + markers["lysogeny_associated_regulators"] = ( + "curated lysogeny regulator or bounded repressor/immunity annotation" + ) + if annot in MAJOR_CAPSID_EXACT: + markers["major_capsid"] = "curated exact major capsid/head/coat annotation" + if re.search(r"\bportal protein\b", annot): + markers["portal"] = "word-bounded annotation match: portal protein" + if re.search(r"\bterminase\b", annot): + markers["terminase"] = "word-bounded annotation match: terminase" + if cat == "tail": + markers["tail"] = "exact PHROG category match: tail" + if re.search(r"\bholin\b", annot): + markers["holin"] = "word-bounded annotation match: holin" + tail_or_virion_lysozyme = re.search( + r"\b(?:virion|tail|baseplate)\b.*\blysozyme\b" + r"|\blysozyme\b.*\b(?:virion|tail|baseplate)\b", + annot, + ) + if annot in ENDOLYSIN_EXACT or re.search(r"\bendolysin\b", annot) or tail_or_virion_lysozyme: + markers["endolysin"] = ( + "curated muralytic annotation, bounded endolysin, or tail/virion " + "lysozyme match" + ) + + if "integrase" in markers or "excisionase" in markers: + markers["integrase_plus_excisionase"] = "union of integrase and excisionase" + if "holin" in markers or "endolysin" in markers: + markers["holin_plus_endolysin"] = "union of holin and endolysin" + return markers + + +def is_known_function(annotation, category): + return ( + normalise_text(category) != "unknown function" + and normalise_text(annotation) not in UNKNOWN_ANNOTATIONS + ) + + def add_count(counts, sample, feature, value): if feature: counts[feature][sample] += value @@ -43,10 +235,11 @@ def format_number(value): return f"{value:.12g}" -def write_matrix(path, features, samples, denominator=None): +def write_matrix(path, features, samples, denominator=None, ordered_features=None): with open(path, "w", encoding="utf-8") as output: output.write("\t" + "\t".join(samples) + "\n") - for feature in sorted(features): + feature_names = ordered_features if ordered_features is not None else sorted(features) + for feature in feature_names: output.write(feature) for sample in samples: value = features[feature].get(sample, 0.0) @@ -57,15 +250,106 @@ def write_matrix(path, features, samples, denominator=None): output.write("\n") +def validate_report_metadata(report, line_number, report_value, database_value, field): + if report_value.strip() and normalise_text(report_value) != normalise_text(database_value): + raise SystemExit( + f"FATAL: {report}:{line_number} {field} {report_value!r} conflicts " + f"with annotation database value {database_value!r}" + ) + + +def process_report_line( + line, + line_number, + report, + sample, + annotation_path, + annotation_database, + total_hits, + phrog_hits, + phrog_counts, + annotation_counts, + category_counts, + marker_counts, + metadata, + observed_database_records, +): + fields = line.rstrip("\r\n").split("\t") + if len(fields) < 2: + raise SystemExit( + f"FATAL: {report}:{line_number} has {len(fields)} columns; expected at least 2" + ) + try: + count = float(fields[1]) + except ValueError: + raise SystemExit( + f"FATAL: {report}:{line_number} has a non-numeric alignment count: " + f"{fields[1]!r}" + ) from None + if not math.isfinite(count) or count < 0: + raise SystemExit( + f"FATAL: {report}:{line_number} has an invalid alignment count: " + f"{fields[1]!r}" + ) + + total_hits[sample] += count + if len(fields) < 12: + return + phrog_id, color, report_annotation, report_category = fields[8:12] + if not phrog_id: + return + try: + canonical_id = normalise_phrog_id(phrog_id) + except ValueError as error: + raise SystemExit(f"FATAL: {report}:{line_number}: {error}") from None + if canonical_id not in annotation_database: + raise SystemExit( + f"FATAL: {report}:{line_number} PHROG {phrog_id!r} is absent from " + f"{annotation_path}" + ) + db_annotation, db_category = annotation_database[canonical_id] + validate_report_metadata( + report, line_number, report_annotation, db_annotation, "annotation" + ) + validate_report_metadata( + report, line_number, report_category, db_category, "category" + ) + + phrog_hits[sample] += count + add_count(phrog_counts, sample, phrog_id, count) + add_count(annotation_counts, sample, report_annotation, count) + add_count(category_counts, sample, report_category, count) + marker_counts["all_phrog_hits"][sample] += count + known_marker = ( + "known_function_phrog_hits" + if is_known_function(db_annotation, db_category) + else "unknown_function_phrog_hits" + ) + marker_counts[known_marker][sample] += count + for marker in classify_phrog(db_annotation, db_category): + marker_counts[marker][sample] += count + + current_metadata = (color, report_annotation, report_category) + if phrog_id in metadata and metadata[phrog_id] != current_metadata: + raise SystemExit( + f"FATAL: {report}:{line_number} has conflicting annotations for {phrog_id}" + ) + metadata[phrog_id] = current_metadata + observed_database_records[canonical_id] = (db_annotation, db_category) + + def main(): args = parse_args() if not os.path.isdir(args.directory): raise SystemExit(f"FATAL: {args.directory} is not a directory") + annotation_database = read_phrog_annotations(args.phrog_annotations) phrog_counts = defaultdict(lambda: defaultdict(float)) annotation_counts = defaultdict(lambda: defaultdict(float)) category_counts = defaultdict(lambda: defaultdict(float)) + marker_counts = defaultdict(lambda: defaultdict(float)) metadata = {} + observed_database_records = {} total_hits = defaultdict(float) phrog_hits = defaultdict(float) samples = [] @@ -74,119 +358,143 @@ def main(): sample_directory = os.path.join(args.directory, sample) if not os.path.isdir(sample_directory): continue - report = os.path.join(sample_directory, sample + REPORT_SUFFIX) if not os.path.isfile(report): if args.verbose: print(f"Skipping {sample}: {report} does not exist", file=sys.stderr) continue - if args.verbose: print(f"Reading {sample} from {report}", file=sys.stderr) samples.append(sample) - with gzip.open(report, "rt", encoding="utf-8") as input_file: - for line_number, line in enumerate(input_file, start=1): - fields = line.rstrip("\r\n").split("\t") - if len(fields) < 2: - raise SystemExit( - f"FATAL: {report}:{line_number} has {len(fields)} columns; " - "expected at least 2" - ) - - try: - count = float(fields[1]) - except ValueError: - raise SystemExit( - f"FATAL: {report}:{line_number} has a non-numeric alignment " - f"count: {fields[1]!r}" - ) from None - if not math.isfinite(count) or count < 0: - raise SystemExit( - f"FATAL: {report}:{line_number} has an invalid alignment " - f"count: {fields[1]!r}" + try: + with gzip.open(report, "rt", encoding="utf-8") as input_file: + for line_number, line in enumerate(input_file, start=1): + process_report_line( + line, + line_number, + report, + sample, + args.phrog_annotations, + annotation_database, + total_hits, + phrog_hits, + phrog_counts, + annotation_counts, + category_counts, + marker_counts, + metadata, + observed_database_records, ) - - total_hits[sample] += count - if len(fields) < 12: - continue - - phrog_id, color, annotation, category = fields[8:12] - if not phrog_id: - continue - - phrog_hits[sample] += count - add_count(phrog_counts, sample, phrog_id, count) - add_count(annotation_counts, sample, annotation, count) - add_count(category_counts, sample, category, count) - - current_metadata = (color, annotation, category) - if phrog_id in metadata and metadata[phrog_id] != current_metadata: - raise SystemExit( - f"FATAL: {report}:{line_number} has conflicting annotations " - f"for {phrog_id}" - ) - metadata[phrog_id] = current_metadata + except (OSError, UnicodeError) as error: + raise SystemExit(f"FATAL: cannot read MMseqs report {report}: {error}") from None if not samples: - raise SystemExit( - f"FATAL: no *{REPORT_SUFFIX} files were found below {args.directory}" - ) + raise SystemExit(f"FATAL: no *{REPORT_SUFFIX} files were found below {args.directory}") os.makedirs(args.output, exist_ok=True) - levels = { - "phrog": phrog_counts, - "annotation": annotation_counts, - "category": category_counts, - } + levels = {"phrog": phrog_counts, "annotation": annotation_counts, "category": category_counts} for level, counts in levels.items(): - write_matrix( - os.path.join(args.output, f"{args.name}_{level}_raw.tsv"), - counts, - samples, - ) - write_matrix( - os.path.join(args.output, f"{args.name}_{level}_norm_all.tsv"), - counts, - samples, - total_hits, - ) - write_matrix( - os.path.join(args.output, f"{args.name}_{level}_norm_phrog.tsv"), - counts, - samples, - phrog_hits, - ) + write_matrix(os.path.join(args.output, f"{args.name}_{level}_raw.tsv"), counts, samples) + write_matrix(os.path.join(args.output, f"{args.name}_{level}_norm_all.tsv"), counts, samples, total_hits) + write_matrix(os.path.join(args.output, f"{args.name}_{level}_norm_phrog.tsv"), counts, samples, phrog_hits) - with open( - os.path.join(args.output, f"{args.name}_phrog_metadata.tsv"), - "w", - encoding="utf-8", - ) as output: + write_matrix(os.path.join(args.output, f"{args.name}_marker_raw.tsv"), marker_counts, samples, ordered_features=MARKER_ORDER) + write_matrix(os.path.join(args.output, f"{args.name}_marker_norm_all.tsv"), marker_counts, samples, total_hits, MARKER_ORDER) + write_matrix(os.path.join(args.output, f"{args.name}_marker_norm_phrog.tsv"), marker_counts, samples, phrog_hits, MARKER_ORDER) + + with open(os.path.join(args.output, f"{args.name}_phrog_metadata.tsv"), "w", encoding="utf-8") as output: output.write("phrog_id\tcolor\tannotation\tcategory\n") for phrog_id in sorted(metadata): output.write(phrog_id + "\t" + "\t".join(metadata[phrog_id]) + "\n") - with open( - os.path.join(args.output, "README.md"), "w", encoding="utf-8" - ) as output: - output.write( - f"""# PHROG count tables + with open(os.path.join(args.output, f"{args.name}_phrog_marker_mapping.tsv"), "w", encoding="utf-8", newline="") as output: + writer = csv.writer(output, delimiter="\t", lineterminator="\n") + writer.writerow(("phrog_id", "annotation", "category", "marker_class", "classification_rule")) + for canonical_id in sorted(observed_database_records, key=int): + annotation, category = observed_database_records[canonical_id] + status = "known_function_phrog_hits" if is_known_function(annotation, category) else "unknown_function_phrog_hits" + relationships = {"all_phrog_hits": "observed MMseqs top hit with a PHROG match", status: "authoritative annotation/category known-function status"} + relationships.update(classify_phrog(annotation, category)) + for marker in MARKER_ORDER: + if marker in relationships: + writer.writerow((f"phrog_{canonical_id}", annotation, category, marker, relationships[marker])) + + with open(os.path.join(args.output, f"{args.name}_phrog_count_summary.tsv"), "w", encoding="utf-8") as output: + output.write("sample\ttotal_mmseqs_top_hits\tall_phrog_hits\tknown_function_phrog_hits\tunknown_function_phrog_hits\tphrog_match_percent\n") + for sample in samples: + all_hits = marker_counts["all_phrog_hits"].get(sample, 0.0) + known = marker_counts["known_function_phrog_hits"].get(sample, 0.0) + unknown = marker_counts["unknown_function_phrog_hits"].get(sample, 0.0) + if not math.isclose(known + unknown, all_hits, rel_tol=1e-12, abs_tol=1e-9): + raise SystemExit(f"FATAL: known and unknown PHROG counts do not sum to all PHROG hits for {sample}") + percent = all_hits * 100 / total_hits[sample] if total_hits[sample] > 0 else 0.0 + output.write( + sample + + "\t" + + "\t".join( + format_number(value) + for value in (total_hits[sample], all_hits, known, unknown, percent) + ) + + "\n" + ) + + with open(os.path.join(args.output, "README.md"), "w", encoding="utf-8") as output: + output.write(f"""# PHROG count tables + +These read-level summaries use the authoritative PHROG annotation database +`{args.phrog_annotations}`. The mapping file includes observed PHROGs only, with +one row per PHROG-marker relationship and the rule that assigned it. The `{args.name}_*_raw.tsv` files contain alignment counts from column 2 of the -MMseqs top-hit reports. - -The `{args.name}_*_norm_all.tsv` files contain counts per million of all MMseqs -top hits, including hits without a PHROG match. - -The `{args.name}_*_norm_phrog.tsv` files contain counts per million of MMseqs -top hits that have a PHROG match. - -Tables are provided by PHROG ID, annotation, and category. The -`{args.name}_phrog_metadata.tsv` file maps each observed PHROG ID to its color, -annotation, and category. -""" - ) +MMseqs top-hit reports. The `{args.name}_*_norm_all.tsv` files contain counts per +1,000,000 all MMseqs top-hit reads, including reads without a PHROG match; this is +the primary normalisation. The `{args.name}_*_norm_phrog.tsv` files contain counts +per 1,000,000 PHROG-matched reads. + +Existing tables are provided by PHROG ID, report annotation, and report category. + +## Marker definitions + +* `all_phrog_hits`: every top hit with a PHROG match. +* `known_function_phrog_hits`: PHROG hits whose authoritative category is not + `unknown function` and whose annotation is not empty or an unknown placeholder. +* `unknown_function_phrog_hits`: the complementary PHROG-matched hits. Known plus + unknown therefore equals all PHROG hits for each sample. +* `integrase`: a word-bounded `integrase` annotation. +* `excisionase`: curated exact excisionase and recombination-directionality-factor + annotations, excluding DNA excision-repair helicases. +* `integrase_plus_excisionase`: the set union of the preceding two classes. +* `lysogeny_associated_regulators`: curated phage repression, immunity, + establishment, and reversal regulators plus bounded repressor/anti-repressor + annotations; generic transcriptional regulators and activators are excluded. +* `major_capsid`: curated exact major capsid, major head, and major coat annotations; + minor and spore coat proteins are excluded. +* `portal`: a word-bounded `portal protein` annotation. +* `terminase`: a word-bounded `terminase` annotation, including large and small + subunits. +* `tail`: the exact PHROG `tail` category. The `connector` category is not merged. +* `holin`: a word-bounded `holin` annotation, including holin/anti-holin. +* `endolysin`: curated endolysin, lysozyme-domain, amidase, and tail/virion muralytic + annotations; bounded endolysin and tail/virion-associated lysozyme descriptions + are also included. Hemolysins and spanins without an independent endolysin + annotation are excluded. +* `holin_plus_endolysin`: the set union of the preceding two classes. + +Combined classes count each read once. Individual component classes can overlap; +for example, a Cox-like excisionase/repressor contributes to both excisionase and +lysogeny-associated regulators, and a tail-associated lysin contributes to tail and +endolysin. + +Marker definitions are annotation-based read summaries, not a definitive lytic- +versus-lysogenic classification and not a contig-level lifestyle analysis. +Structural and lysis genes occur in both temperate and virulent phages, and absence +of an integrase does not prove that a phage is virulent. + +`{args.name}_phrog_metadata.tsv` preserves metadata observed in reports, +`{args.name}_phrog_marker_mapping.tsv` audits marker assignments, and +`{args.name}_phrog_count_summary.tsv` records denominators and PHROG matching QC. +""") if __name__ == "__main__": diff --git a/pawsey_shortread/count_phrog_functions.slurm b/pawsey_shortread/count_phrog_functions.slurm index 62bbd9f..5bbd9fd 100644 --- a/pawsey_shortread/count_phrog_functions.slurm +++ b/pawsey_shortread/count_phrog_functions.slurm @@ -17,7 +17,61 @@ fi source DEFINITIONS.sh : "${SAMPLENAME:?SAMPLENAME must be defined in DEFINITIONS.sh}" +: "${PAWSEY_PROJECT:?PAWSEY_PROJECT must be set}" + +PHROG_DIR=/scratch/$PAWSEY_PROJECT/$USER/Databases/phrog +PHROG_ANNOT="$PHROG_DIR/phrog_annot_v4.tsv.gz" +if [[ ! -s "$PHROG_ANNOT" ]]; then + echo "PHROG annotation database is missing or empty: $PHROG_ANNOT" >&2 + echo "Run mmseqs_add_phrog_function.slurm to download and validate the PHROG database." >&2 + exit 2 +fi +if ! gzip -t "$PHROG_ANNOT"; then + echo "PHROG annotation database is not a valid gzip file: $PHROG_ANNOT" >&2 + echo "Remove or replace the damaged file, then rerun mmseqs_add_phrog_function.slurm." >&2 + exit 2 +fi + +OUTPUT_DIR=phrog_functions +TMP_OUTPUT="${OUTPUT_DIR}.tmp.${SLURM_JOB_ID:-$$}" +if [[ -e "$TMP_OUTPUT" ]]; then + echo "Temporary output path already exists: $TMP_OUTPUT" >&2 + exit 2 +fi +mkdir -- "$TMP_OUTPUT" +PUBLISH_LOCK="${OUTPUT_DIR}.publish.lock" +PUBLISH_LOCK_HELD=0 +cleanup() { + status=$? + if [[ $PUBLISH_LOCK_HELD -eq 1 ]]; then + if ! rmdir -- "$PUBLISH_LOCK"; then + echo "Unable to remove publish lock $PUBLISH_LOCK" >&2 + [[ $status -ne 0 ]] || status=1 + fi + fi + if [[ $status -eq 0 ]]; then + rm -rf -- "$TMP_OUTPUT" + else + echo "Preserving temporary output at $TMP_OUTPUT (exit $status)" >&2 + fi + exit "$status" +} +trap cleanup EXIT export PYTHONPATH="${PYTHONPATH:+${PYTHONPATH}:}$HOME/GitHubs/atavide_lite" -python "$HOME/atavide_lite/bin/count_phrog_functions.py" -d mmseqs -n "$SAMPLENAME" -find phrog_functions -type f -name '*.tsv' -exec pigz {} + +python "$HOME/GitHubs/atavide_lite/bin/count_phrog_functions.py" \ + -d mmseqs -o "$TMP_OUTPUT" -n "$SAMPLENAME" -a "$PHROG_ANNOT" +find "$TMP_OUTPUT" -type f -name '*.tsv' -exec pigz -- {} + + +# mkdir is the exclusive check-and-publish lock: only one concurrent job can +# inspect OUTPUT_DIR and rename its completed temporary directory at a time. +if ! mkdir -- "$PUBLISH_LOCK"; then + echo "Another job is publishing $OUTPUT_DIR; preserving $TMP_OUTPUT" >&2 + exit 2 +fi +PUBLISH_LOCK_HELD=1 +if [[ -e "$OUTPUT_DIR" ]]; then + echo "$OUTPUT_DIR already exists; refusing to overwrite or mix results" >&2 + exit 2 +fi +mv "$TMP_OUTPUT" "$OUTPUT_DIR" diff --git a/tests/test_count_phrog_functions.py b/tests/test_count_phrog_functions.py new file mode 100644 index 0000000..7ac8da1 --- /dev/null +++ b/tests/test_count_phrog_functions.py @@ -0,0 +1,170 @@ +"""Focused tests for read-level PHROG marker counting.""" + +import csv +import gzip +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + + +REPOSITORY = Path(__file__).resolve().parents[1] +SCRIPT = REPOSITORY / "bin" / "count_phrog_functions.py" + +ANNOTATIONS = [ + ("1", "integrase", "integration and excision"), + ("phrog_2", "Cox-like excisionase and repressor", "transcription regulation"), + ("3", "helicase subunit of the DNA excision repair complex", "DNA metabolism"), + ("4", "transcriptional regulator", "transcription regulation"), + ("5", "major capsid protein", "head and packaging"), + ("6", "spore coat protein", "other"), + ("7", "minor coat protein", "head and packaging"), + ("8", "portal protein", "connector"), + ("9", "terminase large subunit", "head and packaging"), + ("10", "tail associated lysin", "tail"), + ("11", "holin/anti-holin", "lysis"), + ("12", "endolysin", "lysis"), + ("13", "hemolysin", "other"), + ("14", "", "unknown function"), + ("15", "recombination directionality factor", "integration and excision"), +] + + +def report_row(count, phrog="", annotation="", category="", color=""): + return ["target", str(count), "x", "x", "x", "x", "x", "x", phrog, color, annotation, category] + + +def read_matrix(path): + with path.open(encoding="utf-8") as handle: + rows = list(csv.reader(handle, delimiter="\t")) + samples = rows[0][1:] + return {row[0]: dict(zip(samples, map(float, row[1:]))) for row in rows[1:]} + + +class CountPhrogFunctionsTest(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.root = Path(self.temporary.name) + self.database = self.root / "phrog.tsv.gz" + with gzip.open(self.database, "wt", encoding="utf-8") as output: + output.write("phrog\tcolor\tannot\tcategory\n") + for phrog, annotation, category in ANNOTATIONS: + output.write(f"{phrog}\tblue\t{annotation}\t{category}\n") + + rows = [ + report_row(2, "phrog_1", "integrase", "integration and excision", "blue"), + report_row(3, "phrog_2", "Cox-like excisionase and repressor", "transcription regulation", "blue"), + report_row(5, "phrog_3", ANNOTATIONS[2][1], "DNA metabolism", "blue"), + report_row(7, "phrog_4", "transcriptional regulator", "transcription regulation", "blue"), + report_row(11, "phrog_5", "major capsid protein", "head and packaging", "blue"), + report_row(13, "phrog_6", "spore coat protein", "other", "blue"), + report_row(17, "phrog_7", "minor coat protein", "head and packaging", "blue"), + report_row(19, "phrog_8", "portal protein", "connector", "blue"), + report_row(23, "phrog_9", "terminase large subunit", "head and packaging", "blue"), + report_row(29, "phrog_10", "tail associated lysin", "tail", "blue"), + report_row(31, "phrog_11", "holin/anti-holin", "lysis", "blue"), + report_row(37, "phrog_12", "endolysin", "lysis", "blue"), + report_row(41, "phrog_13", "hemolysin", "other", "blue"), + report_row(43, "phrog_14", "", "unknown function", "blue"), + report_row(47, "phrog_15", "recombination directionality factor", "integration and excision", "blue"), + report_row(67), + ] + self.write_report("sample_a", rows) + self.write_report("sample_zero", [report_row(100)]) + self.output = self.root / "output" + + def tearDown(self): + self.temporary.cleanup() + + def write_report(self, sample, rows): + directory = self.root / "mmseqs" / sample + directory.mkdir(parents=True) + path = directory / f"{sample}_tophit_report_phrog_function.tsv.gz" + with gzip.open(path, "wt", encoding="utf-8") as output: + for row in rows: + output.write("\t".join(row) + "\n") + + def run_script(self, database=None, expect_success=True): + result = subprocess.run( + [sys.executable, str(SCRIPT), "-d", str(self.root / "mmseqs"), "-o", str(self.output), "-n", "test", "-a", str(database or self.database)], + text=True, + capture_output=True, + ) + if expect_success and result.returncode: + self.fail(result.stderr) + return result + + def test_outputs_counts_normalisation_exclusions_and_zero_sample(self): + self.run_script() + expected_legacy = { + "test_phrog_raw.tsv", "test_phrog_norm_all.tsv", "test_phrog_norm_phrog.tsv", + "test_annotation_raw.tsv", "test_annotation_norm_all.tsv", "test_annotation_norm_phrog.tsv", + "test_category_raw.tsv", "test_category_norm_all.tsv", "test_category_norm_phrog.tsv", + "test_phrog_metadata.tsv", + } + self.assertTrue(expected_legacy.issubset({path.name for path in self.output.iterdir()})) + with (self.output / "test_phrog_raw.tsv").open() as handle: + self.assertEqual(handle.readline().rstrip("\n"), "\tsample_a\tsample_zero") + self.assertEqual(handle.readline().rstrip("\n"), "phrog_1\t2\t0") + + raw = read_matrix(self.output / "test_marker_raw.tsv") + expected = { + "all_phrog_hits": 328, "known_function_phrog_hits": 285, + "unknown_function_phrog_hits": 43, "integrase": 2, "excisionase": 50, + "integrase_plus_excisionase": 52, "lysogeny_associated_regulators": 3, + "major_capsid": 11, "portal": 19, "terminase": 23, "tail": 29, + "holin": 31, "endolysin": 66, "holin_plus_endolysin": 97, + } + for marker, count in expected.items(): + self.assertEqual(raw[marker]["sample_a"], count) + self.assertEqual(raw[marker]["sample_zero"], 0) + + norm_all = read_matrix(self.output / "test_marker_norm_all.tsv") + norm_phrog = read_matrix(self.output / "test_marker_norm_phrog.tsv") + self.assertAlmostEqual( + norm_all["all_phrog_hits"]["sample_a"], 328 * 1_000_000 / 395, places=5 + ) + self.assertEqual(norm_phrog["all_phrog_hits"]["sample_a"], 1_000_000) + self.assertEqual(norm_phrog["all_phrog_hits"]["sample_zero"], 0) + + with (self.output / "test_phrog_count_summary.tsv").open() as handle: + summary = {row["sample"]: row for row in csv.DictReader(handle, delimiter="\t")} + self.assertEqual(float(summary["sample_a"]["known_function_phrog_hits"]) + float(summary["sample_a"]["unknown_function_phrog_hits"]), float(summary["sample_a"]["all_phrog_hits"])) + + mapping = (self.output / "test_phrog_marker_mapping.tsv").read_text() + self.assertIn("phrog_2\tCox-like excisionase and repressor\ttranscription regulation\texcisionase", mapping) + self.assertIn("phrog_2\tCox-like excisionase and repressor\ttranscription regulation\tlysogeny_associated_regulators", mapping) + excluded = ( + "phrog_3\t" + ANNOTATIONS[2][1] + "\tDNA metabolism\texcisionase", + "phrog_4\ttranscriptional regulator\ttranscription regulation\tlysogeny_associated_regulators", + "phrog_6\tspore coat protein\tother\tmajor_capsid", + "phrog_7\tminor coat protein\thead and packaging\tmajor_capsid", + "phrog_13\themolysin\tother\tholin", + "phrog_13\themolysin\tother\tendolysin", + ) + for relationship in excluded: + self.assertNotIn(relationship, mapping) + + def test_uncompressed_annotation_input(self): + plain = self.root / "phrog.tsv" + with gzip.open(self.database, "rt", encoding="utf-8") as source: + plain.write_text(source.read(), encoding="utf-8") + self.run_script(plain) + + def test_malformed_and_conflicting_annotation_records_fail_clearly(self): + malformed = self.root / "malformed.tsv" + malformed.write_text("phrog\tannot\n1\tintegrase\n", encoding="utf-8") + result = self.run_script(malformed, expect_success=False) + self.assertNotEqual(result.returncode, 0) + self.assertIn("must contain tab-separated columns", result.stderr) + + conflicting = self.root / "conflicting.tsv" + conflicting.write_text("phrog\tannot\tcategory\n1\tintegrase\tintegration\nphrog_1\tportal protein\tconnector\n", encoding="utf-8") + result = self.run_script(conflicting, expect_success=False) + self.assertNotEqual(result.returncode, 0) + self.assertIn("conflicting record for PHROG 1", result.stderr) + + +if __name__ == "__main__": + unittest.main()