Latent-state tHMM on CRISPR recorder phylogenies: KP-Tracer analysis (#1014)

[aarmey]

This PR works through #1014, apart from the spatial stretch goals (skipped as requested) and the baseMEMOIR stretch goal (it needs 290 GB of raw images). The full write-up, with figures, is in analysis/kptracer/REPORT.md.

Model changes (lineage/)

• Per-edge transitions: the E-step now takes one transition matrix per edge, exp(Q·t_e). The existing shared-T path is unchanged. Q gets an exact endpoint-conditioned EM update (Hobolth & Jensen), in ctmc.py.
• Edge-wise joint posteriors P(z_p, z_c | Y), plus switch-count, plasticity and heritability summaries.
• Gaussian emission (StateDistributionGaussian): diagonal plus ridge covariance; NaN rows are unobserved, with log-likelihood 0.
• Log-space emissions and a log-space upward pass. The upward pass needed this because KP-Tracer polytomies with thousands of children overflowed np.prod.
• Tree loader (tree_io.py) for Newick, networkx or CassiopeiaTree. It prunes, collapses unifurcations, aligns leaves to AnnData, and computes LCA alleles and mutation-count branch lengths.
• New modules:
  • PhyloHMM (phyloHMM.py) runs EM and reports BIC, held-out likelihood and the lineage-information gain.
  • phylo_sim.py is the simulator.
  • phylo_stats.py holds the Fitch/Sankoff and PATH-style baselines.
• Tests: 19 new tests. They include an exact check against brute-force enumeration on a multifurcating tree, CTMC statistics against numerical quadrature, and a 5,000-child polytomy. The full suite (86 tests), ruff and ty pass.

Findings (details in the report)

• Simulation: switch counts are unbiased, while k-means followed by parsimony over-counts by 1.4–2.9× when states overlap. Ancestral states beat even parsimony given the true labels when only 10% of leaves are sampled. Runtime is linear, about 1.2 s per E-step on a 60k-node tree.
• Tumor-level plasticity: on 21 KP-Tracer sgNT tumors (28,939 cells), tHMM plasticity tracks Fitch parsimony (Spearman ρ = 0.94) and the published scPlasticity (ρ = 0.94).
• States: AT2-like is the stable, most likely root state. High plasticity, AT1-like and the gastric/endoderm states are the least stable, as the paper reports.
• Direction: Early gastric → AT1-like, Late gastric → Endoderm-like and Early EMT-1 → Pre-EMT are directed in at least 92% of tumor bootstraps. AT1 ↔ AT2, which Fitch reports as one-way, comes out bidirectional.
• Lineage information gain: the tree improves held-out prediction by about 1.5 nats per masked cell.
• Caveats:
  • Mutation-count branch lengths are no better than unit lengths.
  • Several states are supported by only one to four tumors; both Mesenchymal states occur only in one nearly star-shaped tree.
  • Unsupervised state discovery is driven by cell density.
• BDMM-Prime on 100 cells: accuracy on the well-determined rates is comparable. The tHMM fits in 0.4 s; BDMM-Prime needs about 7–10 min to reach ESS 200.

Notes for review

• Analysis intermediates (results/*.pkl) are git-ignored. The data come from Zenodo 5847462; the steps are in analysis/kptracer/README.md.
• prepare_data.py needs anndata (uv run --with anndata), so it has a ty ignore on that import. No new dependencies were added.
• The local uv.lock change is not part of this PR.

🤖 Generated with Claude Code

https://claude.ai/code/session_01A7EXGcBoM1ESZnrruqPQff