Repository navigation
Unify likelihood model - #10
Merged
jeromekelleher merged 7 commits intoOct 28, 2026
Merged
Conversation
Remove weight_by_n: recombination is always divided by the number of nodes n, giving one sub-stochastic (killed) Li and Stephens chain, which docs/likelihood_model.md defines. The forward algorithm now starts from 1/n per node, so the log likelihood is P(h, path survives). The weight_by_n=False (n = 1) transition with recombination rho is exactly (1 - rho) / (1 - r) times this chain's transition with r = tsls.model.recombination_probability(rho, n), for every pair of states. So it has the same Viterbi paths and forward-backward posteriors, and only its log likelihood was wrong. sc2ts's model is the r from this conversion; with its parameters (k = 3) that is r = 0.048 per site at n = 20k and 0.89 at n = 3.4M, which is why its posteriors are diffuse. TestModelEquivalence checks the identity on the dense reference. Viterbi depends only on the switch odds q = (r/n) / (1 - r + r/n) and mismatch odds e = mu / (1 - (A - 1) mu): normalised values are exp(-(cost - min cost)) (test_viterbi_values_are_odds). From this, Hmm(prune_threshold=t) drops nodes whose normalised value is below t. That is guaranteed to find an optimal path whose odds are at least t, and tsls.model.path_threshold(k, ...) = e^k q^s (1 - 1e-9) is the largest such threshold for k mismatches and s switches (test_threshold_is_tight). Returned paths with odds >= t are certified optimal (test_certificate). Unlike the experimental pruning option, the prune is applied as each value is normalised, so sparse sites still work. likelihood_threshold stays as an underflow floor. sc2ts's floor at 0.99 mu^k was 6-15% below the exact e^k, had no effect at k = 3 (below q), and as a floor above q at k = 1, 2 could return suboptimal paths.
The model dicts and verification.make_matcher still describe the model as tsinfer-hmm takes it (rho and weight_by_n). The tsls engine and make_matcher convert weight_by_n=False to r = tsls.model.recombination_probability(rho, num_nodes), which has the same Viterbi paths and posteriors. n_weighted_rho is that function. With a mismatch threshold, the tsls benchmark engine now prunes at the exact tsls.model.path_threshold instead of flooring at 0.99 mu^k; tsinfer-hmm keeps the sc2ts floor. Checks: - benchmarking tests: 67 passed, 1 xfailed. - verification test_dense_agreement_random_args (600 queries, half with weight_by_n=False through the conversion against dense with n = 1): forward/backward relative error max 2e-15/5e-15, Viterbi gap max 6e-14. - test_calibration_default and test_calibration_no_weight_by_n: credible set coverage within 0.003 of the set mass at every level (e.g. 0.9827 vs 0.9823 at the 90% level without weight_by_n).
exact_thresholds.py runs the 10 paired.py sc2ts queries on the 100k sub-ARG with no threshold, sc2ts's floor at 0.99 mu^k, and the exact prune at tsls.model.path_threshold(k). Exact: 1.14 s, mean traceback 4761. At k = 0, floor 0.50 s / 979 / 5 of 10 optimal, prune 0.54 s / 608 / 8 of 10. At k = 1, floor 0.83 s / 2537 / 6, prune 0.84 s / 1977 / 9. At k = 2 both are near exact, and at k = 3 neither has any effect (0.99 mu^3 is below q). The certified paths (1, 2, 4, 4 of 10 at k = 0..3) were all optimal under both methods.
Hmm.mean_num_values and max_num_values give the length of the likelihood list after each site (the nodes carrying their own value), and mean_num_updated the number of node values computed per site. For a full site that is the list before coalescing, and for a sparse Viterbi site only the nodes recomputed. The counters cover Viterbi, forward and backward, are reset at each run, and are for the thresholding experiments.
evaluate.py runs each strategy against the exact run, query by query: Viterbi over the whole query (exact; floors; lossless prune; the exact prune for k mismatches; top_n; and escalation from k = 0 until a path is certified), and forward-backward on a 5,000-site window (exact; floors; relative; top_n). It reports the loss in path log odds and certificates, the error in the forward log likelihood and in the posteriors at 50 sites, the nodes tracked and updated per site (the new Hmm counters), and times. Workloads: the sc2ts sub-ARGs, tsinfer chr20p ancestor groups, and sample matching with mu set so that k = 1, 2, 4 mismatches cost one switch, plus tsinfer's mu = 1e-20. Dead end: the first attempt ran forward-backward over the whole chr20p sequence and ran the machine out of memory. A pilot on 5,000-site windows stored about 720 forward and 1,070 backward values per site for sample matching at the k = 4 mu (up to 1,600 for some queries), so all 581k sites would need about 11 GB for backward alone. Forward-backward now runs on windows, and each worker is limited to 5 GB of address space, so a run that runs out of memory is recorded instead of taking the machine down.
…hold The first full run showed a RuntimeWarning in the posterior: with relative or top_n pruning on the tsinfer workloads, forward and backward prune independently and can keep no node in common at a site, so the posterior mass is 0 and the normalised posterior NaN. pandas dropped those sites from the means, hiding the worst failures. Such sites now count as wholly wrong (TV 1) and are reported as no_mass_sites. Add relative_q: drop forward/backward nodes below c q times the maximum, with c = 0.01 and 1, where q is the switch odds. Such a node has less mass than recombination alone gives every node, so this is the forward analogue of the lossless Viterbi prune. In the first run, relative 1e-8 (about 0.4 q on chr20p samples) was the only pruning with errors near 0 on the tsinfer workloads. Also summarise escalation separately: the time ratio, and the k at which each query was certified.
Results of evaluate.py (one full run; 10 sc2ts queries on each of the 10k and 100k sub-ARGs, 6 queries from each of three chr20p ancestor groups, and 8 sample queries at each of four mu, about 20 minutes with 2 workers). Forward-backward runs on 5,000-site windows, because the full chr20p sequence needs about 16 GB at about 2,300 stored values per site. Viterbi: - No certified path was suboptimal (74 certified of 490 runs). - The lossless prune and floors below q never lost anything. - Floors above q gave wrong paths on all sample queries and ancestor group 1485 (log-odds losses of 1e3-1e6). - Exact Viterbi tracks only 2-15 nodes on average on the tsinfer data, and sc2ts sparse sites update 1-15 nodes per site, so pruning saves little time without losing optimality. - Escalation costs 2.2-5.2x exact. Sample paths (100-250 switches) are never certified by the k-mismatch budget. - top_n is 50-270x slower on sc2ts. Forward-backward: - Only thresholds at about the switch odds q or below (relative_q, relative 1e-8, floor 1e-8) stay accurate: logL error <= 0.003 when pruning and posterior TV <= 3e-4, saving 10-28% with mismatches. - Coarser relative pruning loses all posterior mass at up to 86 of 400 sites, because forward and backward prune independently. - Floors >= 1e-4 put the log likelihood off by 1e3-1e4. Also mark the earlier simulated study as superseded.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.