diff --git a/scripts/recombination_params b/scripts/recombination_params new file mode 100755 index 000000000..1c9fbdba2 --- /dev/null +++ b/scripts/recombination_params @@ -0,0 +1,270 @@ +#!/usr/bin/env python3 + +""" +Compute recombination HMM parameters from a dataset's reference tree and write them into pathogen.json. + +This is a Python mirror of the Rust estimator in +`packages/nextclade/src/analyze/recombination_estimate.rs` of the nextclade repository. It must produce +the same values as the Rust run-time fallback, so that a dataset can carry frozen, reviewed parameters: + + gamma = 1 / L + mu_w = mean terminal branch length (substitutions) / L + mu_r = median pairwise inter-clade leaf-to-leaf distance (substitutions) / L + +where L is the reference length. Branch lengths are counted from the per-branch nucleotide mutation +lists in the Auspice tree (`branch_attrs.mutations.nuc`), which avoids the ambiguous Auspice divergence +units. Indels are excluded, matching the Rust model: a token counts only when both its reference and +query bases are present (non-gap), so deletions (query base `-`) and insertions (reference base `-`) +are both skipped and only true substitutions are counted. `mu_r` is the median over all pairs of leaves drawn +from two different clades, where a leaf-to-leaf distance is the number of mutations along the tree path +through the pair's MRCA (`root_distance(a) + root_distance(b) - 2 * root_distance(mrca)`). + +With `--pathogen` the estimates are frozen into that file in place; without it they are printed as JSON +to stdout for inspection. + +Example: + + ./scripts/recombination_params \ + --input-tree data/.../tree.json \ + --input-ref data/.../reference.fasta \ + --pathogen data/.../pathogen.json + + ./scripts/recombination_params \ + --input-tree data/.../tree.json \ + --input-ref data/.../reference.fasta +""" + +from __future__ import annotations + +import argparse +import json +import math +import re +import sys +from collections.abc import Iterator +from dataclasses import dataclass, field +from itertools import combinations +from pathlib import Path +from statistics import median +from typing import Any + +from Bio import SeqIO + +NUC_MUTATIONS_KEY = "nuc" + +# Mirror the Rust `NucSub` parser (`packages/nextclade/src/analyze/nuc_sub.rs`): an unanchored +# `` token whose ref and query are single IUPAC nucleotide letters or the gap `-`. +NUC_MUTATION_REGEX = re.compile(r"([A-Z-])(\d{1,10})([A-Z-])") + +# Valid nucleotide characters accepted by the Rust `to_nuc` (`packages/nextclade/src/alphabet/nuc.rs`): +# the IUPAC codes plus the gap. Note: no 'U' and no lowercase, matching the Rust mapping exactly. +VALID_NUCS = frozenset("ACGTWYMHKRDSBVN-") + +GAP = "-" + + +def main() -> None: + args = parse_args() + + ref_len = read_reference_length(args.input_ref) + root = build_tree(json.loads(args.input_tree.read_text())["tree"]) + params = estimate_params(root, ref_len) + + if params is None: + print( + "Recombination parameters could not be estimated (fewer than two clades or degenerate tree)", + file=sys.stderr, + ) + return + + # With no `--pathogen` target the parameters are the sole output: print them as JSON to stdout so the + # result can be inspected or piped without mutating any file. + if args.pathogen is None: + print(json.dumps(params, indent=2)) + return + + pathogen = json.loads(args.pathogen.read_text()) + recombination = pathogen.get("recombination") or {} + if not isinstance(recombination, dict): + raise ValueError(f'Expected "recombination" in {args.pathogen} to be an object, but got {type(recombination).__name__}') + # Preserve an explicit enablement decision (default on, matching the run-time default); freeze the estimates. + recombination.setdefault("enabled", True) + recombination.update(params) + pathogen["recombination"] = recombination + args.pathogen.write_text(json.dumps(pathogen, indent=2) + "\n") + + print(f"Wrote recombination parameters to {args.pathogen}: {params}") + + +def estimate_params(root: Node, ref_len: int) -> dict[str, float] | None: + # An empty or header-only reference leaves the model unresolved rather than dividing by zero. + if ref_len <= 0: + return None + + gamma = as_probability(1.0 / ref_len) + mu_w = estimate_mu_w(root, ref_len) + mu_r = estimate_mu_r(root, ref_len) + + # Mirror `RecombinationHmmParams::new`: every rate must be a valid probability, the chain must be + # sticky (`gamma < 0.5`, so state switching is rarer than staying), and the recombinant rate must be + # elevated (`mu_r > mu_w`), otherwise the two HMM states are indistinguishable. Any violation leaves + # the model unresolved, matching the run-time skip. + if gamma is None or gamma >= 0.5 or mu_w is None or mu_r is None or mu_r <= mu_w: + return None + + return {"gamma": gamma, "muW": mu_w, "muR": mu_r} + + +def estimate_mu_w(root: "Node", ref_len: int) -> float | None: + branch_lengths = [node.branch_muts for node in iter_nodes(root) if not node.children] + if not branch_lengths: + return None + return as_probability((sum(branch_lengths) / len(branch_lengths)) / ref_len) + + +def estimate_mu_r(root: "Node", ref_len: int) -> float | None: + # Group by LEAF clade only: `mu_r` is the median pairwise distance between leaves of different + # clades, so a clade that labels only internal nodes cannot contribute (matching the Rust + # `clade_leaves` grouping and the leaf-derived "fewer than two clades" gate). + clade_leaves: dict[str, list[Node]] = {} + for node in iter_nodes(root): + if not node.children and node.clade is not None: + clade_leaves.setdefault(node.clade, []).append(node) + + if len(clade_leaves) < 2: + return None + + # Precompute root distances for ALL nodes so each leaf pair costs one lookup per endpoint plus a + # single MRCA search. + root_dists = {node: root_distance(node) for node in iter_nodes(root)} + + # Exhaustive pairwise distances between leaves of different clades. + distances: list[float] = [] + for leaves_a, leaves_b in combinations(clade_leaves.values(), 2): + for a in leaves_a: + # `a` is fixed across the inner loop, so compute its ancestor set once and reuse it. + ancestors_a = set(ancestors_inclusive(a)) + for b in leaves_b: + mrca = mrca_with_ancestors(ancestors_a, b) + distances.append(float(root_dists[a] + root_dists[b] - 2 * root_dists[mrca])) + + if not distances: + return None + + return as_probability(median(distances) / ref_len) + + +def root_distance(node: "Node") -> int: + total = 0 + current: Node | None = node + while current is not None: + total += current.branch_muts + current = current.parent + return total + + +def mrca_with_ancestors(ancestors_a: set["Node"], b: "Node") -> "Node": + # Walk `b` toward the root until it reaches a node in `a`'s inclusive ancestor set. + current: Node | None = b + while current is not None: + if current in ancestors_a: + return current + current = current.parent + raise ValueError("Nodes have no common ancestor") + + +def ancestors_inclusive(node: "Node") -> list["Node"]: + ancestors = [] + current: Node | None = node + while current is not None: + ancestors.append(current) + current = current.parent + return ancestors + + +def iter_nodes(root: Node) -> Iterator[Node]: + stack = [root] + while stack: + node = stack.pop() + yield node + stack.extend(node.children) + + +def build_tree(auspice_root: dict[str, Any]) -> Node: + # Build iteratively: real datasets can be deeply ladderized, and recursion depth would then track + # the number of tips and overflow the interpreter's stack (matching the iterative Rust builder). + root = node_from_auspice(auspice_root, None) + stack: list[tuple[dict[str, Any], Node]] = [(auspice_root, root)] + while stack: + auspice_node, node = stack.pop() + for auspice_child in auspice_node.get("children", []): + child = node_from_auspice(auspice_child, node) + node.children.append(child) + stack.append((auspice_child, child)) + return root + + +def node_from_auspice(auspice_node: dict[str, Any], parent: Node | None) -> Node: + muts = auspice_node.get("branch_attrs", {}).get("mutations", {}).get(NUC_MUTATIONS_KEY, []) + return Node( + clade=auspice_node.get("node_attrs", {}).get("clade_membership", {}).get("value"), + branch_muts=count_branch_substitutions(muts), + parent=parent, + ) + + +def count_branch_substitutions(muts: list[str]) -> int: + # Number of true nucleotide substitutions on a branch (mirrors the Rust `count_nuc_mutations`). A + # token counts only when both its reference and query bases are present (non-gap). Indels are + # excluded because the decoder never observes them: deletions ("A15-", gap query) are treated as + # missing data, and insertions ("-10A", gap reference) have no observation channel at all, so + # counting either would calibrate the emission rates against events the HMM cannot score. A token + # that does not parse is a malformed tree annotation and raises rather than being silently + # miscounted. Nextclade-built trees carry substitutions only; indels appear on externally produced + # trees (augur/TreeTime). + count = 0 + for m in muts: + match = NUC_MUTATION_REGEX.search(m) + if match is None or match.group(1) not in VALID_NUCS or match.group(3) not in VALID_NUCS: + raise ValueError(f"Unable to parse nucleotide mutation: '{m}'") + if match.group(1) != GAP and match.group(3) != GAP: + count += 1 + return count + + +def as_probability(x: float) -> float | None: + return x if (math.isfinite(x) and 0.0 < x < 1.0) else None + + +def read_reference_length(path: Path) -> int: + records = list(SeqIO.parse(path, "fasta")) + if len(records) != 1: + raise ValueError(f"Expected exactly one sequence record in {path}, but found {len(records)}") + return len(records[0].seq) + + +# eq=False gives identity-based equality and hashing, which is required because nodes reference their +# parent (structural equality would recurse) and are stored in sets keyed by identity. +@dataclass(eq=False) +class Node: + clade: str | None + branch_muts: int + parent: "Node | None" = None + children: list["Node"] = field(default_factory=list) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, description=__doc__) + parser.add_argument("--input-tree", type=Path, required=True, help="Path to the reference tree.json") + parser.add_argument("--input-ref", type=Path, required=True, help="Path to the reference.fasta") + parser.add_argument( + "--pathogen", + type=Path, + default=None, + help="Path to pathogen.json to modify in place; if omitted, the estimated parameters are printed as JSON to stdout", + ) + return parser.parse_args() + + +if __name__ == "__main__": + main() diff --git a/tests/test_recombination_params.py b/tests/test_recombination_params.py new file mode 100644 index 000000000..d20e82d3e --- /dev/null +++ b/tests/test_recombination_params.py @@ -0,0 +1,309 @@ +"""Cross-implementation parity tests for the recombination parameter estimator. + +The authoritative implementation is the Rust estimator in +`packages/nextclade/src/analyze/recombination_estimate.rs` of the nextclade repository. Each fixture and +expected value here is transcribed from that file's unit tests, which serve as the oracle: the Python +helper must produce bit-identical parameters so a dataset can carry frozen, reviewed values. + +The estimator does deterministic integer arithmetic (small counts divided by ref_len), so the results +are exact; the tight `delta` compares them to the Rust literals essentially exactly. +""" + +import json +import subprocess +import sys +import unittest +from importlib.machinery import SourceFileLoader +from importlib.util import module_from_spec, spec_from_loader +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import Any + +# The helper is an extensionless executable, so load it by explicit path rather than a normal import. +_SCRIPT_PATH = Path(__file__).resolve().parent.parent / "scripts" / "recombination_params" +_loader = SourceFileLoader("recombination_params", str(_SCRIPT_PATH)) +_spec = spec_from_loader(_loader.name, _loader) +assert _spec is not None +rp = module_from_spec(_spec) +# Register before executing so `@dataclass` can resolve the `Node` forward reference (dataclasses look +# the module up in `sys.modules` to evaluate string annotations under `from __future__ import annotations`). +sys.modules[_loader.name] = rp +_loader.exec_module(rp) + +# Reference length used by every Rust fixture; keeps the expected literals identical to the oracle. +REF_LEN = 100 + +# Tightest tolerance that passes for the exact integer/ref_len arithmetic. +DELTA = 1e-12 + + +def _leaf(name: str, clade: str, nuc: list[str]) -> dict[str, Any]: + return { + "name": name, + "node_attrs": {"clade_membership": {"value": clade}}, + "branch_attrs": {"mutations": {"nuc": nuc}}, + } + + +def _params(tree: dict[str, Any]) -> dict[str, float] | None: + return rp.estimate_params(rp.build_tree(tree), REF_LEN) + + +class TestRecombinationParamsParity(unittest.TestCase): + def _assert_params(self, params: dict[str, float] | None, gamma: float, mu_w: float, mu_r: float) -> None: + assert params is not None, "expected resolved parameters, got None" + self.assertAlmostEqual(gamma, params["gamma"], delta=DELTA) + self.assertAlmostEqual(mu_w, params["muW"], delta=DELTA) + self.assertAlmostEqual(mu_r, params["muR"], delta=DELTA) + + def test_recombination_params_all_params_from_two_clade_tree(self) -> None: + # Rust `test_recombination_estimate_all_params_from_tree`: gamma 0.01, mu_w 0.02, mu_r 0.05. + # Leaves A1 (rd 2), B1 (rd 2), B2 (rd 4); mean terminal (2+1+3)/3 = 2. Inter-clade leaf pairs + # A1<->B1 = 4, A1<->B2 = 6; median{4, 6} = 5. + tree = { + "name": "root", + "children": [ + _leaf("A1", "A", ["A10C", "A20G"]), + { + "name": "B", + "node_attrs": {"clade_membership": {"value": "B"}}, + "branch_attrs": {"mutations": {"nuc": ["A30T"]}}, + "children": [ + _leaf("B1", "B", ["A40C"]), + _leaf("B2", "B", ["A50C", "A60G", "A70T"]), + ], + }, + ], + } + self._assert_params(_params(tree), gamma=1.0 / 100.0, mu_w=2.0 / 100.0, mu_r=5.0 / 100.0) + + def test_recombination_params_three_clades_uses_nonroot_mrca(self) -> None: + # Rust `test_recombination_estimate_three_clades_uses_nonroot_mrca`: mu_w 0.04, mu_r 0.09. + # A1 and B1 share the non-root MRCA I, exercising the `- 2 * root_distance(mrca)` term. + # Root distances A1=3, B1=5, C1=6; leaf pairs {6, 9, 11}, median 9. Terminals 2, 4, 6 -> mean 4. + tree = { + "name": "root", + "children": [ + { + "name": "I", + "node_attrs": {}, + "branch_attrs": {"mutations": {"nuc": ["A5C"]}}, + "children": [ + _leaf("A1", "A", ["A10C", "A20G"]), + _leaf("B1", "B", ["A30T", "A40C", "A50G", "A60T"]), + ], + }, + _leaf("C1", "C", ["A70C", "A80G", "A90T", "A100C", "A110G", "A120T"]), + ], + } + self._assert_params(_params(tree), gamma=1.0 / 100.0, mu_w=4.0 / 100.0, mu_r=9.0 / 100.0) + + def test_recombination_params_nested_clades_uses_leaf_distance(self) -> None: + # Rust `test_recombination_estimate_nested_clades_uses_leaf_distance`: mu_w 0.04, mu_r 0.09. + # Clade ancestors sit one mutation apart, but mu_r reflects inter-clade LEAF divergence. + # Leaf pairs sorted {6, 6, 8, 8, 10, 10, 12, 14} -> median (8+10)/2 = 9. Terminals 1,8,6,1,4 -> 4. + tree = { + "name": "root", + "children": [ + { + "name": "P", + "node_attrs": {"clade_membership": {"value": "parent"}}, + "branch_attrs": {"mutations": {"nuc": ["A1C"]}}, + "children": [ + _leaf("P1", "parent", ["A2C"]), + { + "name": "C", + "node_attrs": {"clade_membership": {"value": "child"}}, + "branch_attrs": {"mutations": {"nuc": ["A3C"]}}, + "children": [ + _leaf("C1", "child", ["A10C", "A11C", "A12C", "A13C", "A14C", "A15C", "A16C", "A17C"]), + _leaf("C2", "child", ["A20C", "A21C", "A22C", "A23C", "A24C", "A25C"]), + ], + }, + _leaf("P2", "parent", ["A30C"]), + ], + }, + _leaf("O1", "other", ["A40C", "A41C", "A42C", "A43C"]), + ], + } + self._assert_params(_params(tree), gamma=1.0 / 100.0, mu_w=4.0 / 100.0, mu_r=9.0 / 100.0) + + def test_recombination_params_excludes_deletions_from_branch_length(self) -> None: + # Rust `test_recombination_estimate_excludes_deletions_from_branch_length`: the `A15-` deletion is + # excluded, so B1's terminal branch is 1 (substitution A40C only); mean (2+1)/2 = 1.5 -> mu_w 0.015. + tree = { + "name": "root", + "children": [ + _leaf("A1", "A", ["A10C", "A20G"]), + _leaf("B1", "B", ["A40C", "A15-"]), + ], + } + params = _params(tree) + assert params is not None + self.assertAlmostEqual(1.5 / 100.0, params["muW"], delta=DELTA) + + def test_recombination_params_excludes_insertions_from_branch_length(self) -> None: + # Rust `test_recombination_estimate_excludes_insertions`: the `-10A` insertion (gap reference) is + # an indel the decoder never observes, so it is excluded; B1's terminal branch is 1 (substitution + # A40C only); mean (2+1)/2 = 1.5 -> mu_w 0.015. + tree = { + "name": "root", + "children": [ + _leaf("A1", "A", ["A10C", "A20G"]), + _leaf("B1", "B", ["A40C", "-10A"]), + ], + } + params = _params(tree) + assert params is not None + self.assertAlmostEqual(1.5 / 100.0, params["muW"], delta=DELTA) + + def test_recombination_params_single_clade_is_unresolved(self) -> None: + # Rust `test_recombination_estimate_single_clade_is_unresolved`: mu_r undefined with one clade. + tree = { + "name": "root", + "children": [ + _leaf("L1", "A", ["A10C"]), + _leaf("L2", "A", ["A20G"]), + ], + } + self.assertIsNone(_params(tree)) + + def test_recombination_params_internal_only_clade_is_unresolved(self) -> None: + # Rust `test_recombination_estimate_internal_only_clade_is_fewer_than_two_clades`: internal node I + # is labeled clade "B" but no leaf is, so only one LEAF clade ("A") exists and mu_r is undefined. + tree = { + "name": "root", + "children": [ + { + "name": "I", + "node_attrs": {"clade_membership": {"value": "B"}}, + "branch_attrs": {"mutations": {"nuc": ["A5C"]}}, + "children": [ + _leaf("A1", "A", ["A10C"]), + _leaf("A2", "A", ["A20G"]), + ], + }, + ], + } + self.assertIsNone(_params(tree)) + + def test_recombination_params_no_branch_mutations_is_unresolved(self) -> None: + # Rust `test_recombination_estimate_no_branch_mutations_skip_reason`: two clades but no `nuc` + # branch mutations, so every rate is zero and the model is left unresolved. + tree = { + "name": "root", + "children": [ + {"name": "A1", "node_attrs": {"clade_membership": {"value": "A"}}}, + {"name": "B1", "node_attrs": {"clade_membership": {"value": "B"}}}, + ], + } + self.assertIsNone(_params(tree)) + + def test_recombination_params_errors_on_malformed_branch_mutation(self) -> None: + # Rust `test_recombination_estimate_errors_on_malformed_branch_mutation`: a malformed annotation + # surfaces as an error instead of being silently miscounted. + tree = { + "name": "root", + "children": [ + _leaf("A1", "A", ["not-a-mutation"]), + _leaf("B1", "B", ["A40C"]), + ], + } + with self.assertRaises(ValueError) as ctx: + _params(tree) + self.assertEqual("Unable to parse nucleotide mutation: 'not-a-mutation'", str(ctx.exception)) + + +class TestCountBranchSubstitutions(unittest.TestCase): + def test_count_branch_substitutions_counts_substitutions(self) -> None: + self.assertEqual(2, rp.count_branch_substitutions(["A10C", "A20G"])) + + def test_count_branch_substitutions_empty_is_zero(self) -> None: + self.assertEqual(0, rp.count_branch_substitutions([])) + + def test_count_branch_substitutions_excludes_deletion(self) -> None: + # `A15-` has a gap query base (deletion) and is not counted; `A40C` is. + self.assertEqual(1, rp.count_branch_substitutions(["A40C", "A15-"])) + + def test_count_branch_substitutions_excludes_insertion(self) -> None: + # `-10A` has a gap reference base (insertion) and is not counted; only the substitution `A40C` is. + self.assertEqual(1, rp.count_branch_substitutions(["A40C", "-10A"])) + + def test_count_branch_substitutions_counts_ambiguous_iupac(self) -> None: + # `N` and `R` are valid IUPAC nucleotide codes accepted by the Rust `to_nuc`. + self.assertEqual(1, rp.count_branch_substitutions(["N5R"])) + + def test_count_branch_substitutions_raises_on_unparseable(self) -> None: + with self.assertRaises(ValueError): + rp.count_branch_substitutions(["not-a-mutation"]) + + def test_count_branch_substitutions_raises_on_invalid_nucleotide(self) -> None: + # `Z` matches the position regex `[A-Z-]` but is not a nucleotide, so `to_nuc` (mirrored by the + # VALID_NUCS check) rejects it. + with self.assertRaises(ValueError): + rp.count_branch_substitutions(["A10Z"]) + + +# The two-clade oracle tree (gamma 0.01, mu_w 0.02, mu_r 0.05 at ref_len 100) as an Auspice file body. +_TWO_CLADE_AUSPICE = { + "version": "v2", + "meta": {}, + "tree": { + "name": "root", + "children": [ + _leaf("A1", "A", ["A10C", "A20G"]), + { + "name": "B", + "node_attrs": {"clade_membership": {"value": "B"}}, + "branch_attrs": {"mutations": {"nuc": ["A30T"]}}, + "children": [ + _leaf("B1", "B", ["A40C"]), + _leaf("B2", "B", ["A50C", "A60G", "A70T"]), + ], + }, + ], + }, +} + + +class TestRecombinationParamsCli(unittest.TestCase): + """End-to-end CLI smoke tests over the full argument-parsing and I/O path.""" + + def _write_inputs(self, tmp: Path) -> tuple[Path, Path]: + tree = tmp / "tree.json" + tree.write_text(json.dumps(_TWO_CLADE_AUSPICE)) + ref = tmp / "reference.fasta" + ref.write_text(">ref\n" + "A" * REF_LEN + "\n") + return tree, ref + + def _run(self, *args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, str(_SCRIPT_PATH), *args], + capture_output=True, + text=True, + check=True, + ) + + def test_recombination_params_cli_prints_json_when_no_pathogen(self) -> None: + with TemporaryDirectory() as tmp_dir: + tree, ref = self._write_inputs(Path(tmp_dir)) + result = self._run("--input-tree", str(tree), "--input-ref", str(ref)) + params = json.loads(result.stdout) + self.assertAlmostEqual(1.0 / 100.0, params["gamma"], delta=DELTA) + self.assertAlmostEqual(2.0 / 100.0, params["muW"], delta=DELTA) + self.assertAlmostEqual(5.0 / 100.0, params["muR"], delta=DELTA) + + def test_recombination_params_cli_writes_pathogen_in_place(self) -> None: + with TemporaryDirectory() as tmp_dir: + tmp = Path(tmp_dir) + tree, ref = self._write_inputs(tmp) + pathogen = tmp / "pathogen.json" + pathogen.write_text(json.dumps({"schemaVersion": "3.0.0"})) + self._run("--input-tree", str(tree), "--input-ref", str(ref), "--pathogen", str(pathogen)) + written = json.loads(pathogen.read_text())["recombination"] + self.assertEqual(True, written["enabled"]) + self.assertAlmostEqual(5.0 / 100.0, written["muR"], delta=DELTA) + + +if __name__ == "__main__": + unittest.main()