Repository navigation
Compression improvements - #9
Merged
jeromekelleher merged 6 commits intoOct 26, 2026
Merged
Conversation
paired.py only timed the tsinfer chr20p workloads, and required tsls to find exactly tsinfer-hmm's path. It now also times three sc2ts workloads (num_mismatches = 3, weight_by_n=False), on spread queries from data/queries_scaling: - sc2ts_n100k_k1: 10 queries on the 100k sub-ARG with sc2ts's k = 1 likelihood threshold. sc2ts matches every sample at k = 0 and k = 1 first, and reruns those whose cost exceeds k + 1 at high precision. - sc2ts_n100k: the same queries, exact (the high precision reruns). - sc2ts_n300k: 4 queries on the 300k sub-ARG, exact. --tsinfer and --sc2ts select either set. A path that differs from tsinfer-hmm's is now classed by its log likelihood under the model (engines.model_path_log_likelihood, which takes per-site rates, n and per-site numbers of alleles) as a tie, better or worse, and the table reports these counts and both engines' traceback memory. With a likelihood threshold, ties in the floored values can differ in model likelihood, so differences there are expected (see the docstring). Baseline (3868de8), tsls / tsinfer-hmm wall time, two full runs: | workload | queries | ratio | tsls memory (MiB) | |---|---|---|---| | ancestors_g84 | 50 | 0.76, 0.75 | 64 | | ancestors_g1000 | 50 | 0.82, 0.74 | 64 | | ancestors_g1483 | 100 | 0.91, 0.76 | 64 | | ancestors_g1484 | 50 | 0.91, 0.74 | 64 | | ancestors_g1485 | 100 | 0.98, 0.90 | 64 | | samples | 6 | 1.14, 0.80 | 64 | | sc2ts_n100k_k1 | 10 | 1.05, 1.02 | 256 | | sc2ts_n100k | 10 | 1.09, 1.06 | 960 | | sc2ts_n300k | 4 | 1.00, 0.95 | 1792 | tsls is level with or slower than tsinfer-hmm on sc2ts. The tsinfer ratios vary more between these runs than in the README's (samples 1.14 in the first run, during which other jobs ran on the machine); the later builds were run alternately with this one, and compared within rounds. perf on two sc2ts_n100k queries put about 40% of find_path in the walk up from each likelihood node to find its allele, and 40% in the coalescing walks; on a tsinfer sample, the allele walk was 58%. Instrumenting the engine showed that on sc2ts only 5-6.5 of 670-9,200 stored values (1.6-1.8 of 370-2,400 with k = 1) change at each site, bit for bit, and that the maximum likelihood node mismatches the query at only 2-11 sites per query.
At each site, the Viterbi, forward and backward loops found the allele of every node in the likelihood list by walking up the tree to the first node with a mutation at the site, or a root. Walks from different nodes share most of their paths, so the cost was the sum of the path lengths: on sc2ts, where the list holds thousands of samples under deep chains of ancestors, and on tsinfer's deep trees, this was the largest single cost (about 40% of find_path on sc2ts and 58% on a tsinfer sample). get_memoised_allelic_state now gives every node on the path the state it finds, so later walks stop where an earlier one passed, and records these nodes so they are unset after the loop. The cost is the size of the union of the paths. Only nodes without a state are written, so roots and mutation nodes keep their marks. On sc2ts the walk fell to about 18% of find_path (perf), and coalescing became the largest cost. Paired wall time, tsls / tsinfer-hmm, two full runs alternating with the other builds (base is 99e5c3d's baseline): | workload | base | memoised | |---|---|---| | ancestors_g84 | 0.76, 0.75 | 0.78, 0.78 | | ancestors_g1000 | 0.82, 0.74 | 0.71, 0.71 | | ancestors_g1483 | 0.91, 0.76 | 0.81, 0.68 | | ancestors_g1484 | 0.91, 0.74 | 0.75, 0.68 | | ancestors_g1485 | 0.98, 0.90 | 0.78, 0.69 | | samples | 1.14, 0.80 | 0.76, 0.57 | | sc2ts_n100k_k1 | 1.05, 1.02 | 0.87, 0.85 | | sc2ts_n100k | 1.09, 1.06 | 0.88, 0.86 | | sc2ts_n300k | 1.00, 0.95 | 0.79, 0.73 | 10-30% faster everywhere except the earliest group (4 ms queries), with identical paths and memory.
The traceback stored the whole (node, recombination_required) list at every site, unless it was identical to the previous site's. On sc2ts the lists hold hundreds to thousands of nodes but only a few change per site (experiments/compression estimated 50-430x less storage for a delta encoding), and full-ARG queries used 3-7 GB. Each site now stores a traceback_delta_t: the nodes whose flag changed, or that entered or left the list, with their flags at the previous site. Each node's flag as of the last site stored is kept in the spare byte of tsls_node_t (traceback_flag), with a copy of the stored list to find the nodes that left it. The traceback starts from the flags as stored at the last site, reads them directly in its walk up the tree, and undoes each site's delta after processing it, which replaces setting and unsetting the whole list at every site and leaves every flag unset at the end. mean_traceback_size still reports the mean list size. The traceback allocator's minimum block is now 4 MiB rather than 64 MiB. Hmm.get_traceback is removed: it was unused, and returned an empty dict after find_path, which clears the traceback of the window. Paired wall time, tsls / tsinfer-hmm, and tsls's traceback memory, two full runs alternating with the other builds: | workload | memoised | + delta | memory (MiB), before -> after | |---|---|---|---| | ancestors_g84 | 0.78, 0.78 | 0.78, 0.77 | 64 -> 4 | | ancestors_g1000 | 0.71, 0.71 | 0.72, 0.70 | 64 -> 4 | | ancestors_g1483 | 0.81, 0.68 | 0.73, 0.68 | 64 -> 4 | | ancestors_g1484 | 0.75, 0.68 | 0.71, 0.67 | 64 -> 8 | | ancestors_g1485 | 0.78, 0.69 | 0.84, 0.84 | 64 -> 4 | | samples | 0.76, 0.57 | 0.55, 0.55 | 64 -> 12 | | sc2ts_n100k_k1 | 0.87, 0.85 | 0.84, 0.81 | 256 -> 4 | | sc2ts_n100k | 0.88, 0.86 | 0.82, 0.80 | 960 -> 4 | | sc2ts_n300k | 0.79, 0.73 | 0.69, 0.69 | 1792 -> 4 | Memory falls to the allocator's minimum block nearly everywhere (from 1.8 GB on sc2ts 300k). Time improves 4-10% on sc2ts and is level on tsinfer within the noise (group 1485 moved the other way in both runs; samples improved in one). Paths are identical to tsinfer-hmm's.
On sc2ts only about 6 of hundreds to thousands of stored likelihoods change at each site (see 99e5c3d), yet every site walked, normalised and coalesced the whole list. In the normalised frame a node's value only changes where it mismatches the query: with uniform rates, once the previous site's maximum likelihood node (value 1) matches, every matching node goes through the same map as at the previous site. tsls_hmm_process_site_sparse uses this. When a tree spans enough sites (remaining sites x list size >= number of nodes), its nodes are numbered in preorder, so each mutation's subtree is a range of positions, and the likelihood list is kept sorted by preorder. At a site where the rates are the same as at the previous site, the previous site was normalised the same way, and the previous maximum likelihood node matches the query, it: - finds the list nodes in mismatching ranges by binary search, and gives a likelihood to each mutation node whose match differs from its parent's; - recomputes the values of those nodes and of the nodes that changed or mismatched at the previous site, with the same arithmetic as a full site, so values are bit-for-bit the same; - coalesces and stores traceback deltas for those nodes only. Other sites (tree changes, the maximum node mismatching, more than 64 mutations, rate changes, viterbi_matrix, pruning, count_distinct) are processed in full, recording which nodes changed when the next site can be sparse. A sparse site keeps the previous maximum likelihood node rather than the first in list order, and nodes below a recomputed node that become redundant stay in the list until the next full site, so paths may differ from processing every site in full only among equally likely ones. Hmm(sparse=False) processes every site in full; num_sparse_sites counts the sparse sites of the last run. The traceback store now takes the nodes that may have changed (the list and the nodes removed from it at the site, or a sparse site's nodes) instead of comparing against a copy of the stored list. Tests: TestSparse compares sparse and full runs, and the dense optimum, on Wright-Fisher ARGs with polytomies, unary nodes, internal samples and recurrent mutations, with and without thresholds, over windows and with changing rates; valgrind memcheck is clean. Two existing tests that expect identical paths from find_path and viterbi_matrix (or pruning) now compare with sparse=False. Paired wall time, tsls / tsinfer-hmm, two full runs alternating with the other builds, with tsls's paths against tsinfer-hmm's (same/tie/better/worse): | workload | delta (ce0334a) | sparse | paths | |---|---|---|---| | sc2ts_n100k_k1 | 0.84, 0.81 | 0.017, 0.017 | 8/1/0/1 | | sc2ts_n100k | 0.82, 0.80 | 0.012, 0.012 | 7/3/0/0 | | sc2ts_n300k | 0.69, 0.69 | 0.051, 0.048 | 3/1/0/0 | The 10 exact 100k queries take 0.75 s against 63 s for tsinfer-hmm. The one "worse" k = 1 path is a tie in the floored values: at site 28,651 of query 88, every node with value 1 had mismatched at the site before, and after renormalising nine nodes tie at value 1, including nodes with 8 and 17 mismatches, since the floor erases how many mismatches a node has made. Full processing (and tsinfer-hmm) takes the first in its list order; the sparse run's list is in preorder. Dead end: the first version recorded changed nodes at every full site, which cost tsinfer 5-15% (samples 0.62, 0.64 against 0.55, 0.55 for delta), although its trees change about every second site and are never numbered. Changes are now only recorded in numbered trees. tsinfer workloads only, two runs alternating with delta: | workload | delta | sparse | |---|---|---| | ancestors_g84 | 0.72, 0.77 | 0.78, 0.75 | | ancestors_g1000 | 0.74, 0.71 | 0.73, 0.74 | | ancestors_g1483 | 0.70, 0.73 | 0.73, 0.72 | | ancestors_g1484 | 0.67, 0.84 | 0.73, 0.69 | | ancestors_g1485 | 0.71, 0.84 | 0.89, 0.72 | | samples | 0.57, 0.68 | 0.58, 0.56 | Level within the noise, with identical paths. (The sc2ts rows above are from the first version; the fix doesn't affect numbered trees.) hmm.c grows by about 500 lines, mostly the sparse site and its helpers.
At k = 1, tsls takes 0.26-0.51 of tsinfer-hmm's time on three queries of the full ARG, with 26 MiB of traceback memory against up to 2.4 GB, but only 4-14% of sites are sparse: trees span about 58 sites, and numbering a whole tree of 3.4M nodes rarely pays. Dead end: numbering trees 16 times more eagerly (remaining sites x list size x 16 >= nodes) made more sites sparse (queries 11 and 22: 6,490 and 9,780 of 29,565) but took longer (35.2 and 40.3 s against 30.6 and 31.5 s), so it was not kept.
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.