From 4d2a67b15ae58440b610332ba14464b4a158f762 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 22 Oct 2026 20:44:48 +0100 Subject: [PATCH 1/2] Coalesce Viterbi likelihoods before storing the traceback Add late-stage tsinfer workloads (1000 Genomes chr20p ancestor groups 1483-1485 and sample matching) to the compression experiments, and measure the proposed delta and affine encodings on them. On these workloads a quarter to a third of the stored Viterbi entries were redundant: when edges change, nodes get explicit copies of their inherited values, and the engine stored the traceback before coalescing them away. That churn limited a delta encoding of the traceback to 1.1-1.7x. Viterbi now coalesces before storing the traceback, and no longer after. A node is removed only if its value and its recombination_required flag both equal the ones it inherits, so the traceback, in which unlisted nodes take the flag of their nearest listed ancestor, is unchanged. Nodes kept for their flag are removed at a later site. Forward and backward still coalesce after storing, on values only. What was tried: 1. Flag-aware coalescing before storing plus the old pass after: minimal lists, but two passes per site made tsls 1.36-1.48x tsinfer-hmm. 2. A single flag-aware pass before storing, with an uncached walk to the nearest listed ancestor: +2-4% instructions over the old engine. 3. As 2, with the inherited flag cached along tree paths alongside the inherited likelihood (recombination_required_cache): landed. It uses no more instructions than the old engine (callgrind: -0.2% on 20 group 1485 ancestors, -0.3% on one sample), and head-to-head wall times agree within noise. Results: the stored lists now equal the change points, Viterbi churn falls 5-8x, engine traceback entries fall 31-47%, and a delta encoding on top saves a further 5-8x. tsls still finds tsinfer's exact path on all benchmark queries. The compression experiments gain: - a tsinfer case in capture.py, with --cases and merged manifests; - churn, value change, flag change and affine residual measures in analysis.py; - an Encodings section and work plots in report.py. The README records the findings and the sequence of attempts. --- experiments/compression/README.md | 158 ++++++++++++ experiments/compression/analysis.py | 145 ++++++++++- experiments/compression/capture.py | 324 +++++++++++++++++++++---- experiments/compression/pyproject.toml | 3 +- experiments/compression/report.py | 189 ++++++++++++++- experiments/compression/uv.lock | 16 +- tsls/lib/matcher.c | 78 +++++- tsls/lib/tsls.h | 3 + 8 files changed, 854 insertions(+), 62 deletions(-) diff --git a/experiments/compression/README.md b/experiments/compression/README.md index 7883ccf..f4d1185 100644 --- a/experiments/compression/README.md +++ b/experiments/compression/README.md @@ -11,6 +11,10 @@ uv run python capture.py --sizes 10000,30000,100000 # about 35 minutes uv run python report.py # results/report.md ``` +`capture.py --cases` captures a subset of the cases (`msprime`, `sc2ts`, +`tsinfer`), replacing their rows in an existing manifest and keeping the +others. `report.py --cases` analyses a subset, for quick iteration. + `capture.py` uses `AncestorMatcher.viterbi_matrix` and `forward_matrix`. It saves each result as a `CompressedMatrix` with `CompressedMatrix.dump`, next to the tree sequence it refers to. Reload a matrix with @@ -30,6 +34,21 @@ The cases captured are: recombinants with one or two breakpoints. - **msprime ARGs**: held-out samples matched against 50 or 1000 other samples, with the same HMM settings. +- **tsinfer workloads**: the late ancestor matching groups 1483, 1484 and + 1485, and sample matching, from the tsinfer 1.0 inference of 1000 Genomes + chr20p (580,909 sites). These need the workloads made by + `benchmarking/tsinfer_data.py` (see `benchmarking/README.md`). They use + tsinfer's settings (ρ = 0.01, μ = 1e-20, `weight_by_n=True`, two + alleles). For each group there are 4 ancestors at spread quantiles of + window length, and there are 4 random sample haplotypes. Each path is + checked against the one tsinfer found. + - Viterbi runs over each query's window: the whole genome for samples. + - `forward_matrix` has no start and end, so forward runs on the tree + sequence restricted to the window's sites (`window_ts`), which keeps the + node IDs and n. For samples the forward window is the middle 20,000 + sites (`--sample-forward-sites`). `capture.py --quick` checks the + windowing by comparing windowed Viterbi with Viterbi over + `[start, end)`. `analysis.py` computes the measures: - From the stored lists, at every site: the list size, distinct values, @@ -42,6 +61,16 @@ The cases captured are: - change points after rounding the values; - an approximate minimum if only samples had to be represented, as in tskit (10 sites). +- Encoding measures, comparing each site with the previous one: + - churn: nodes entering or leaving the stored list; + - changed values of nodes in both lists (a delta encoding); + - affine residuals: values of nodes in both lists that the site's global + map does not predict. For Viterbi the map is + `f(v) = max(v·stay, ρ/n)·e/c`, with `stay = 1 − ρ + ρ/n`, `e` the + matching emission and `c` the normalisation factor; this is exact. For + forward it is `f(v) = a·v + b`, with `a = (1 − ρ)·e/c` and `b` fitted. + - the same on the minimal (change point) lists at the decoded sites, + whose previous sites are decoded too. ## The representation @@ -87,6 +116,11 @@ the site that end up equal to their parent. Representing only the samples, as tskit does, would save at most 5–30% on sc2ts. So the all-nodes model is not what makes the lists large. +This does not hold on tsinfer workloads, where the trees change much more +often. There, a quarter to a third of the Viterbi entries were redundant +until the engine was changed to coalesce before storing. See +[Affine and delta encodings on tsinfer workloads](#affine-and-delta-encodings-on-tsinfer-workloads). + ### 2. Why Viterbi stores so many nodes with so few values on sc2ts The Viterbi values are almost entirely four exact classes: @@ -155,3 +189,127 @@ where `floor = 1e-4` cut forward storage by about 200×. approximation. - The captured matrices (`results/*.npz`) can be used to size and check any new encoding offline before it is implemented in C. + +## Affine and delta encodings on tsinfer workloads + +Means over 4 queries per workload, for late ancestor matching (groups +1483–1485) and sample matching on 1000 Genomes chr20p. Viterbi covers each +query's whole window; forward covers ancestor windows and 20,000 sites for +samples. + +### Before coalescing the traceback + +| | ancestors g1483 | g1484 | g1485 | samples | +|---|---|---|---|---| +| Viterbi stored / site | 12.2 | 11.5 | 20.1 | 26.9 | +| Viterbi change points / site | 9.1 | 8.3 | 13.2 | 19.8 | +| Viterbi churn (entering + leaving) / site | 6.0 | 6.3 | 13.8 | 16.1 | +| Viterbi values changed / site | 0.73 | 0.76 | 1.13 | 1.25 | +| Viterbi affine residuals / site | 0.22 | 0.19 | 0.32 | 0.28 | +| Traceback: engine ÷ delta encoding | 1.6× | 1.6× | 1.1× | 1.7× | +| Forward stored / site | 73 | 80 | 139 | 185 | +| Forward values changed / site | 70 | 77 | 132 | 176 | +| Forward affine residuals / site | 1.6 | 1.6 | 2.8 | 3.1 | +| Forward: stored ÷ affine + delta | 7.8× | 8.5× | 8.0× | 8.5× | + +- Viterbi lists are small (12–27 per site, against hundreds to thousands + on sc2ts), and their values hardly change: about one per site, and only + 0.2–0.3 per site that the global map `max(v·stay, ρ/n)·e/c` doesn't + predict. `weight_by_n=True` doesn't change this, because the uniform stay + factor cancels in normalisation. Exact and tolerant comparisons agree, so + rounding doesn't break a delta encoding. +- But the lists churn: half or more of the entries enter or leave at each + site, because the tree changes about every second site. So delta + encoding the traceback saved only 1.1–1.7×, against 50–430× on sc2ts. +- Most of the churn was the engine's bookkeeping, not structure in the + likelihoods. When edges change, nodes get explicit copies of their + inherited values, and the engine stored the traceback before removing + values equal to their parent's. A quarter to a third of the stored + entries were redundant (7 of 27 for samples), against almost none on + sc2ts. On the + minimal lists the churn was only 1.3–2.2 per site. +- Forward values all change at every site, so delta encoding doesn't help + (0.95×), but the affine map `a·v + b` predicts all but 1.6–3.1 values per + site. Forward has the same redundant entries (3–9 per site), because it + also stores before coalescing. + +### Coalescing before storing the traceback + +The fix is to coalesce before storing. The catch is the traceback: it +stores only `(node, recombination_required)`, and during traceback a node +not in the list takes the flag of its nearest ancestor in the list. A node +with the same value as its parent but a different flag can't be removed +without changing the path. So the coalescing pass, when run before storing +in Viterbi, removes a node only if both its value and its flag equal the +ones it inherits. Nodes kept because of their flag are removed at a later +site, once their flag is recomputed. + +What was tried, in order: + +1. **Flag-aware coalescing before storing, plus the old pass after + storing.** The stored lists became exactly the change points, and the + paths were unchanged. But it ran two coalescing passes per site, and tsls + went from about 1.1× to 1.36–1.48× tsinfer-hmm's time. +2. **Flag-aware coalescing before storing only (one pass per site).** The + stored lists are the same as in 1. The nearest-ancestor walk for the + flag cost 2–4% more instructions than the old engine (table below). Its + time relative to tsinfer-hmm, measured in separate runs, looked like + 1.13–1.22× (against 1.09–1.10× before). This turned out to be mostly + noise in comparing ratios across runs. +3. **As 2, with the inherited flag cached** (landed). The flag is cached + along tree paths with the inherited likelihood + (`ancestor_matcher_get_inherited_likelihood_and_flag`, which uses the new + `recombination_required_cache` array), so the walk is shared. It costs + no more than the old engine. + +Instruction counts, from callgrind with collection limited to +`ancestor_matcher_find_path`, so they are deterministic: + +| | old engine | 2: flag lookup uncached | 3: flag lookup cached | +|---|---|---|---| +| 20 ancestors, group 1485 | 9.92 G | 10.13 G (+2.2%) | 9.89 G (−0.2%) | +| 1 sample | 12.54 G | 13.01 G (+3.8%) | 12.50 G (−0.3%) | + +Head-to-head wall times of the old engine and 3 agree within noise +(4 samples: 11.2 vs 11.3 s; 100 group 1485 ancestors: 14.1 vs 14.4 s, best +of several runs, in alternating order). The 1.36–1.48× for version 1 is +also a ratio to tsinfer-hmm from a separate run. It is well outside the +noise, but not as precise as the counts above. + +The matrices below were captured with version 1; versions 2 and 3 store +the same lists. + +| | ancestors g1483 | g1484 | g1485 | samples | +|---|---|---|---|---| +| Viterbi stored / site | 9.2 | 8.3 | 13.1 | 18.8 | +| Viterbi churn / site | 1.3 | 1.3 | 1.8 | 2.2 | +| Viterbi values changed / site | 0.23 | 0.21 | 0.32 | 0.26 | +| Engine traceback entries (change) | −35% | −35% | −47% | −31% | +| Traceback: engine ÷ delta encoding | 5.1× | 5.0× | 5.2× | 8.4× | +| Traceback: old engine ÷ delta encoding | 7.8× | 7.7× | 9.7× | 12× | + +- The stored lists are now exactly the change points, and churn falls about + 5–8×. Paths are unchanged: tsls still finds tsinfer's path for every + query in all six benchmark workloads (2,014 queries). +- The engine's traceback shrinks by a third to a half, and a delta encoding + on top of it saves another 5–8×. In the benchmarks, mean traceback size + for samples fell from 21.6 to 13.4, and the matcher's memory from 91 to + 64–71 MiB. +- Not changed: forward and backward still coalesce after storing (values + only; they have no flags), so their stored lists keep the redundant + entries described above. +- The sc2ts and msprime results earlier in this README were captured with + the old engine. They had almost no redundant entries, so they shouldn't + change much. + +### Implications for tsinfer + +- Coalesce before storing (done for Viterbi) before trying a delta encoding. + Delta encoding then cuts traceback memory by 5–8×, though tsinfer's + traceback is already small (13 entries per site for samples). +- In an affine Viterbi engine, the per-site work would be dominated by + churn from edge insertions and removals (1–2 per site), not by value + updates (0.1–0.3 per site). +- For forward and backward, the affine update is the clear win: about 8× + fewer explicit updates on tsinfer workloads, and 50–100× on sc2ts. + Forward should also coalesce before storing. diff --git a/experiments/compression/analysis.py b/experiments/compression/analysis.py index 0684ed2..db72b5c 100644 --- a/experiments/compression/analysis.py +++ b/experiments/compression/analysis.py @@ -21,7 +21,45 @@ def site_slices(cm): return [slice(cm.offsets[j], cm.offsets[j + 1]) for j in range(cm.num_sites)] -def list_measures(cm): +class SiteMap: + """ + The global map that the engine applies to the value of every node whose + allele matches the query (or where the query is missing), from one site + to the next. Values that the map does not predict need an explicit + update in an engine that applies the map implicitly (an "affine" + encoding). + + - Viterbi: f(v) = max(v * stay, rho / n) * e / c, where stay is + 1 - rho + rho / n, e is the matching emission probability, and c is + the site's normalisation factor (the maximum). This is exact. + - Forward: f(v) = a * v + b, where a = (1 - rho) * e / c. The shift b is + rho / n * T * e / c, where T is the total at the previous site, which + is not stored, so we fit b as the median over the carried nodes. + + Here n is the number of nodes if weight_by_n, else 1. + """ + + def __init__(self, algorithm, *, rho, mu, n, num_alleles): + self.algorithm = algorithm + self.rho = rho + self.n = n + self.emission = 1 - (num_alleles - 1) * mu + + def predict(self, previous, normalisation_factor, current): + """ + Return the predicted values at a site for nodes with the specified + previous values. For forward, the shift is fitted to the specified + current values. + """ + rho, n, e, c = self.rho, self.n, self.emission, normalisation_factor + if self.algorithm == "viterbi": + return np.maximum(previous * (1 - rho + rho / n), rho / n) * e / c + a = (1 - rho) * e / c + b = np.median(current - a * previous) if len(previous) > 0 else 0 + return a * previous + b + + +def list_measures(cm, site_map=None): """ Return a DataFrame of per-site measures computed from the stored lists: @@ -38,6 +76,20 @@ def list_measures(cm): (Viterbi only). The engine reuses the previous site's list only if the nodes and flags are identical, so traceback memory is the sum of stored over sites with new_list. + + The encoding measures compare each site with the previous one: + + - entered, left: nodes stored at this site and not the previous one, and + the reverse (the churn in the list's membership); + - changed: nodes stored at both sites whose value changed (changed_tol: + by a relative difference of more than 1e-9); + - flag_changed: nodes stored at both sites whose recombination_required + flag changed (Viterbi only); + - affine_residual: nodes stored at both sites whose value is not + predicted by the site's global map (see SiteMap), to a relative + tolerance of 1e-9, if site_map is specified; + - affine_fit: the fraction of nodes stored at both sites whose value is + predicted by the map. """ ts = cm.ts mutation_nodes = [set() for _ in range(ts.num_sites)] @@ -47,13 +99,15 @@ def list_measures(cm): previous_pairs = set() previous_nodes = set() previous_slice = None + previous_nonempty = None for j, s in enumerate(site_slices(cm)): + row = carried_measures(cm, previous_nonempty, s, j, site_map) nodes = cm.nodes[s].tolist() values = cm.values[s].tolist() pairs = set(zip(nodes, values)) node_set = set(nodes) union = len(node_set | previous_nodes) - row = { + row |= { "site": j, "stored": len(nodes), "distinct": len(set(values)), @@ -75,9 +129,40 @@ def list_measures(cm): previous_pairs = pairs previous_nodes = node_set previous_slice = s + if len(nodes) > 0: + previous_nonempty = s return pd.DataFrame(rows) +def carried_measures(cm, previous, current, site, site_map): + """ + Return the dict of encoding measures comparing the stored lists in the + specified previous and current slices (see list_measures). Sites with + empty lists (outside the matched window) are skipped over. + """ + nodes = cm.nodes[current] + if previous is None or len(nodes) == 0: + return {"entered": len(nodes), "left": 0, "changed": 0, "changed_tol": 0} + previous_nodes = cm.nodes[previous] + common, i, k = np.intersect1d(previous_nodes, nodes, return_indices=True) + before = cm.values[previous][i] + after = cm.values[current][k] + row = { + "entered": len(nodes) - len(common), + "left": len(previous_nodes) - len(common), + "changed": int(np.sum(before != after)), + "changed_tol": int(np.sum(~np.isclose(before, after, rtol=1e-9, atol=0))), + } + if cm.flags is not None: + row["flag_changed"] = int(np.sum(cm.flags[previous][i] != cm.flags[current][k])) + if site_map is not None: + predicted = site_map.predict(before, cm.normalisation_factor[site], after) + on_map = np.isclose(predicted, after, rtol=1e-9, atol=0) + row["affine_residual"] = int(np.sum(~on_map)) + row["affine_fit"] = float(np.mean(on_map)) if len(common) > 0 else np.nan + return row + + def iter_trees_at_sites(ts, sites): """ Yield (tree, site) for each of the specified sorted site IDs. @@ -160,7 +245,7 @@ def samples_only_parsimony(tree, labels, is_sample): return num_changes -def decoded_measures(cm, num_sites=200, parsimony_sites=20, seed=1): +def decoded_measures(cm, num_sites=200, parsimony_sites=20, seed=1, site_map=None): """ Return a DataFrame of per-site measures computed by decoding the values at a sample of sites: @@ -173,7 +258,18 @@ def decoded_measures(cm, num_sites=200, parsimony_sites=20, seed=1): - quantised_: change points after rounding log10 values to 1/d; - mean_clade_size, median_clade_size: subtree sizes of the change points; - samples_only: change points if only samples had to be represented (at - a smaller sample of sites). + a smaller sample of sites, and only if the tree sequence has samples). + + The engine stores its lists before removing values equal to the parent's, + and gives nodes their own values when edges change, so the stored lists + can change more between sites than the values do. To measure the + encodings on the minimal lists, we also decode the previous non-empty + site and compare the change points (if site_map is specified): + + - cp_entered, cp_left: nodes that became or stopped being change points; + - cp_changed: change points at both sites whose value changed; + - cp_affine_residual: change points at both sites whose value is not + predicted by the site's global map. """ ts = cm.ts rng = np.random.default_rng(seed) @@ -184,20 +280,34 @@ def decoded_measures(cm, num_sites=200, parsimony_sites=20, seed=1): parsimony = set( rng.choice(sites, size=min(parsimony_sites, len(sites)), replace=False) ) + sampled = set(sites.tolist()) + if site_map is not None: + index = np.searchsorted(nonempty, sites) + previous = nonempty[index[index > 0] - 1] + sites = np.union1d(sites, previous) is_sample = np.zeros(ts.num_nodes, dtype=bool) is_sample[ts.samples()] = True mutation_nodes = {} for site, node in zip(ts.mutations_site, ts.mutations_node): mutation_nodes.setdefault(int(site), set()).add(int(node)) rows = [] + last = None for tree, site in iter_trees_at_sites(ts, sites): value, in_tree = decode_site(cm, tree, site) parent = tree.parent_array cp = change_points(value, in_tree, parent) + cp_nodes = np.where(cp)[0] + current = (site, cp_nodes, value[cp_nodes]) + previous_site = last + last = current + if site not in sampled: + continue s = slice(cm.offsets[site], cm.offsets[site + 1]) stored = cm.nodes[s] redundant = stored[~cp[stored]] - clade_sizes = np.array([tree.num_samples(u) for u in np.where(cp)[0]]) + clade_sizes = np.array([]) + if ts.num_samples > 0: + clade_sizes = np.array([tree.num_samples(u) for u in np.where(cp)[0]]) row = { "site": site, "stored": len(stored), @@ -218,13 +328,36 @@ def decoded_measures(cm, num_sites=200, parsimony_sites=20, seed=1): for digits in [1, 2, 4]: q = np.where(in_tree, quantise(value, digits), np.nan) row[f"quantised_{digits}"] = int(np.sum(change_points(q, in_tree, parent))) - if site in parsimony: + if site in parsimony and ts.num_samples > 0: _, labels = np.unique(np.nan_to_num(value, nan=-1), return_inverse=True) row["samples_only"] = samples_only_parsimony(tree, labels, is_sample) + if site_map is not None and previous_site is not None: + row |= change_point_churn(cm, previous_site, current, site_map) rows.append(row) return pd.DataFrame(rows) +def change_point_churn(cm, previous, current, site_map): + """ + Return the dict of encoding measures comparing the change points + (site, nodes, values) at the previous and current sites. + """ + _, previous_nodes, previous_values = previous + site, nodes, values = current + common, i, k = np.intersect1d(previous_nodes, nodes, return_indices=True) + before = previous_values[i] + after = values[k] + predicted = site_map.predict(before, cm.normalisation_factor[site], after) + return { + "cp_entered": len(nodes) - len(common), + "cp_left": len(previous_nodes) - len(common), + "cp_changed": int(np.sum(before != after)), + "cp_affine_residual": int( + np.sum(~np.isclose(predicted, after, rtol=1e-9, atol=0)) + ), + } + + def value_histogram(cm, top=10): """ Return a DataFrame of the most common stored values and the fraction of diff --git a/experiments/compression/capture.py b/experiments/compression/capture.py index 1a1e61b..b54cafa 100644 --- a/experiments/compression/capture.py +++ b/experiments/compression/capture.py @@ -1,13 +1,17 @@ """ Capture the exact per-site Viterbi and forward values (as tsls -CompressedMatrix files) for queries on sc2ts sub-ARGs and simulated ARGs, so -that compression can be analysed and prototyped offline. +CompressedMatrix files) for queries on sc2ts sub-ARGs, simulated ARGs and +tsinfer workloads, so that compression can be analysed and prototyped +offline. Writes results/.trees (the tree sequence the matrices refer to), -results/__.npz and results/manifest.csv. +results/__.npz and results/manifest.csv. Capturing +some of the cases replaces their rows in an existing manifest and keeps the +others. """ import argparse +import json import pathlib import sys @@ -17,12 +21,17 @@ ROOT = pathlib.Path(__file__).resolve().parents[2] sys.path.insert(0, str(ROOT / "experiments" / "thresholding")) +sys.path.insert(0, str(ROOT / "benchmarking")) +import queries import thresholding as th import tsls v = th.v RESULTS = pathlib.Path(__file__).resolve().parent / "results" SC2TS_DATA = ROOT / "benchmarking" / "data" +TSINFER_DATA = ROOT / "benchmarking" / "data" / "tgp_chr20p" +TSINFER_WORKLOADS = ["ancestors_g1483", "ancestors_g1484", "ancestors_g1485", "samples"] +CASES = ["msprime", "sc2ts", "tsinfer"] def sc2ts_queries(): @@ -49,11 +58,11 @@ def sc2ts_queries(): return chosen -def capture(name, ts, queries, mu, rho, weight_by_n, rows): +def capture(name, ts, query_list, mu, rho, weight_by_n, rows): ts_path = RESULTS / f"{name}.trees" ts.dump(ts_path) mi = tsls.MatcherIndexes(ts, alleles=v.ALLELES) - for query_id, label, h in queries: + for query_id, label, h in query_list: am = v.make_matcher( ts, mu, rho, weight_by_n=weight_by_n, count_distinct=True, matcher_indexes=mi ) @@ -76,6 +85,7 @@ def capture(name, ts, queries, mu, rho, weight_by_n, rows): "mu": mu, "rho": rho, "weight_by_n": weight_by_n, + "num_alleles": len(v.ALLELES), "num_segments": len(parent), "num_mismatches": int(np.sum((matched != h) & (h != -1))), "mean_traceback_size": am.mean_traceback_size, @@ -86,55 +96,285 @@ def capture(name, ts, queries, mu, rho, weight_by_n, rows): print(f"Captured {name} query {query_id} ({label})", flush=True) +def drop_root_edges(ts): + """ + Return the tree sequence with the edges of node 0 removed, so that the + ultimate root of a tsinfer or sc2ts ARG is isolated and not a state. + """ + tables = ts.dump_tables() + tables.edges.keep_rows(tables.edges.parent != 0) + return tables.tree_sequence() + + +def window_ts(ts, start, end): + """ + Return the tree sequence restricted to the sites from start to end + (exclusive). Nodes are kept, so node IDs and the number of nodes (n with + weight_by_n) are unchanged, and running over all its sites is the same + as running over [start, end) of the original. + """ + position = ts.sites_position + right = position[end] if end < ts.num_sites else ts.sequence_length + tables = ts.dump_tables() + keep = np.zeros(ts.num_sites, dtype=bool) + keep[start:end] = True + tables.delete_sites(np.where(~keep)[0]) + tables.keep_intervals([[position[start], right]], simplify=False) + return tables.tree_sequence() + + +def choose_tsinfer_queries(df, kind, num_queries, seed): + """ + Return the query IDs to capture from a tsinfer workload: ancestors at + spread quantiles of window length, and random sample haplotypes. + """ + if kind == "ancestor": + quantiles = [0.1, 0.4, 0.7, 0.95][:num_queries] + order = np.argsort((df.end - df.start).to_numpy(), kind="stable") + index = np.round(np.array(quantiles) * (len(df) - 1)).astype(int) + return [int(df.query_id.iloc[order[j]]) for j in index] + rng = np.random.default_rng(seed) + return sorted(int(q) for q in rng.choice(df.query_id, num_queries, replace=False)) + + +def format_path(left, parent): + return queries.format_list(parent), queries.format_list(left[1:]) + + +def check_window(mi, ts, h, start, end, vm, rho, mu, path): + """ + Check that Viterbi on the windowed tree sequence gives the same path and + stored values as running over [start, end) of the full tree sequence, + which validates the windowing used for forward. + """ + wts = window_ts(ts, start, end) + am = tsls.AncestorMatcher( + tsls.MatcherIndexes(wts), + np.full(wts.num_sites, rho), + np.full(wts.num_sites, mu), + ) + (left, _, parent), _, wvm = am.viterbi_matrix(h[start:end]) + assert np.array_equal(parent, path[2]) + assert np.array_equal(left[:-1], path[0][:-1]) + s = slice(vm.offsets[start], vm.offsets[end]) + assert np.array_equal(wvm.nodes, vm.nodes[s]) + assert np.array_equal(wvm.values, vm.values[s]) + print(" windowed Viterbi matches the full tree sequence", flush=True) + + +def capture_tsinfer(workload, *, num_queries, forward_sites, seed, check, rows): + """ + Capture Viterbi over each chosen query's window, and forward over the + same window (ancestors) or a fixed window in the middle (samples), with + the parameters tsinfer used. + """ + query_dir = TSINFER_DATA / workload + with open(query_dir / "model.json") as f: + model = json.load(f) + assert model["weight_by_n"] and model["alleles"] is None + mu, rho = model["mu"], model["rho"] + H, _, df = queries.load_queries(query_dir) + kind = df.kind.iloc[0] + name = f"tsinfer_{workload}" + ts = drop_root_edges(tskit.load(query_dir / "reference.trees")) + ts_path = RESULTS / f"{name}.trees" + ts.dump(ts_path) + mi = tsls.MatcherIndexes(ts) + for query_id in choose_tsinfer_queries(df, kind, num_queries, seed): + query = df.iloc[query_id] + assert query.query_id == query_id + h = H[query_id] + start, end = queries.query_window(H, query_id) + am = tsls.AncestorMatcher( + mi, np.full(ts.num_sites, rho), np.full(ts.num_sites, mu) + ) + path, matched, vm = am.viterbi_matrix(h, start, end) + left, _, parent = path + truth = (str(query.truth_parents), str(query.truth_breakpoints)) + truth = tuple("" if x == "nan" else x for x in truth) + if format_path(left[::-1], parent[::-1]) != truth: + raise ValueError(f"{name} query {query_id}: path differs from tsinfer") + if check: + check_window(mi, ts, h, start, end, vm, rho, mu, path) + common = { + "arg": name, + "case": "tsinfer", + "group": int(query.group), + "query": query_id, + "label": f"{kind}_{query.sample_id}", + "mu": mu, + "rho": rho, + "weight_by_n": True, + "num_alleles": 2, + "num_segments": len(parent), + "num_mismatches": int(np.sum((matched != h) & (h != -1))), + } + matrix = RESULTS / f"{name}_{query_id}_viterbi.npz" + vm.dump(matrix) + rows.append( + { + **common, + "trees": ts_path.name, + "algorithm": "viterbi", + "matrix": matrix.name, + "num_nodes": ts.num_nodes, + "num_sites": ts.num_sites, + "num_samples": ts.num_samples, + "start": start, + "end": end, + "mean_traceback_size": am.mean_traceback_size, + } + ) + + if kind == "sample": + middle = (start + end) // 2 + fstart = max(start, middle - forward_sites // 2) + fend = min(end, fstart + forward_sites) + else: + fstart, fend = start, end + wts = window_ts(ts, fstart, fend) + wts_path = RESULTS / f"{name}_{query_id}_window.trees" + wts.dump(wts_path) + fam = tsls.AncestorMatcher( + tsls.MatcherIndexes(wts), + np.full(wts.num_sites, rho), + np.full(wts.num_sites, mu), + ) + fm = fam.forward_matrix(h[fstart:fend]) + matrix = RESULTS / f"{name}_{query_id}_forward.npz" + fm.dump(matrix) + rows.append( + { + **common, + "trees": wts_path.name, + "algorithm": "forward", + "matrix": matrix.name, + "num_nodes": wts.num_nodes, + "num_sites": wts.num_sites, + "num_samples": wts.num_samples, + "start": fstart, + "end": fend, + } + ) + np.save(RESULTS / f"{name}_{query_id}_query.npy", h) + print( + f"Captured {name} query {query_id} ({kind}, {end - start} sites, " + f"forward over {fend - fstart})", + flush=True, + ) + + +def case_of(arg): + return arg.split("_")[0] + + +def write_manifest(rows, cases): + """ + Write the manifest, replacing the rows of the specified cases in any + existing manifest. + """ + path = RESULTS / "manifest.csv" + df = pd.DataFrame(rows) + if path.exists(): + old = pd.read_csv(path) + old["case"] = old.arg.map(case_of) + df = pd.concat([old[~old.case.isin(cases)], df], ignore_index=True) + df["case"] = df.arg.map(case_of) + df.to_csv(path, index=False) + + def main(): parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument( + "--cases", + default=",".join(CASES), + help=f"Cases to capture (comma separated, from {', '.join(CASES)})", + ) parser.add_argument( "--sizes", default="10000,30000", help="sc2ts sub-ARG sizes (comma separated)" ) - parser.add_argument("--quick", action="store_true", help="One query of each kind") + parser.add_argument( + "--quick", + action="store_true", + help="One query of each kind, and check the forward windowing", + ) + parser.add_argument( + "--tsinfer-queries", + type=int, + default=4, + help="Queries per tsinfer workload", + ) + parser.add_argument( + "--sample-forward-sites", + type=int, + default=20000, + help="Sites in the forward window for tsinfer samples", + ) args = parser.parse_args() + cases = args.cases.split(",") + for case in cases: + if case not in CASES: + parser.error(f"Unknown case {case!r}") RESULTS.mkdir(exist_ok=True) rows = [] mu, rho = v.hmm_params(3) - # Simulated ARGs: held-out samples, as in the verification suite - for num_samples in [50, 1000]: - for seed in [1] if args.quick else [1, 2]: - ts, h, _, _ = th.threading_case(num_samples, seed) - capture( - f"msprime_n{num_samples}_s{seed}", - ts, - [(0, "held_out", h)], - mu, - rho, - False, - rows, - ) + if "msprime" in cases: + # Simulated ARGs: held-out samples, as in the verification suite + for num_samples in [50, 1000]: + for seed in [1] if args.quick else [1, 2]: + ts, h, _, _ = th.threading_case(num_samples, seed) + capture( + f"msprime_n{num_samples}_s{seed}", + ts, + [(0, "held_out", h)], + mu, + rho, + False, + rows, + ) + + if "sc2ts" in cases: + if (SC2TS_DATA / "subargs").exists(): + H = np.load(SC2TS_DATA / "queries_scaling" / "queries.npz")["H"] + query_list = sc2ts_queries() + if args.quick: + seen = set() + query_list = [ + q for q in query_list if not (q[1] in seen or seen.add(q[1])) + ] + for size in [int(x) for x in args.sizes.split(",")]: + ts = tskit.load(SC2TS_DATA / "subargs" / f"sc2ts_bench_n{size}.trees") + capture( + f"sc2ts_n{size}", + drop_root_edges(ts), + [(query_id, label, H[query_id]) for query_id, label in query_list], + mu, + rho, + False, + rows, + ) + else: + print(f"No sc2ts data in {SC2TS_DATA}: skipping") - # sc2ts sub-ARGs - if (SC2TS_DATA / "subargs").exists(): - H = np.load(SC2TS_DATA / "queries_scaling" / "queries.npz")["H"] - queries = sc2ts_queries() - if args.quick: - seen = set() - queries = [q for q in queries if not (q[1] in seen or seen.add(q[1]))] - for size in [int(x) for x in args.sizes.split(",")]: - ts = tskit.load(SC2TS_DATA / "subargs" / f"sc2ts_bench_n{size}.trees") - tables = ts.dump_tables() - tables.edges.keep_rows(tables.edges.parent != 0) - ts = tables.tree_sequence() - capture( - f"sc2ts_n{size}", - ts, - [(query_id, label, H[query_id]) for query_id, label in queries], - mu, - rho, - False, - rows, + if "tsinfer" in cases: + missing = [w for w in TSINFER_WORKLOADS if not (TSINFER_DATA / w).exists()] + if len(missing) > 0: + print( + f"Missing tsinfer workloads {missing} in {TSINFER_DATA} (see " + "benchmarking/README.md): skipping" ) - else: - print(f"No sc2ts data in {SC2TS_DATA}: skipping") - pd.DataFrame(rows).to_csv(RESULTS / "manifest.csv", index=False) + else: + for workload in TSINFER_WORKLOADS: + capture_tsinfer( + workload, + num_queries=1 if args.quick else args.tsinfer_queries, + forward_sites=args.sample_forward_sites, + seed=1, + check=args.quick, + rows=rows, + ) + write_manifest(rows, cases) if __name__ == "__main__": diff --git a/experiments/compression/pyproject.toml b/experiments/compression/pyproject.toml index b0c477f..10695e3 100644 --- a/experiments/compression/pyproject.toml +++ b/experiments/compression/pyproject.toml @@ -4,6 +4,7 @@ version = "0.0.0" description = "Analysis of likelihood compression in the tsls engine." requires-python = ">=3.11" dependencies = [ + "click>=8.5.0", "matplotlib", "msprime", "numpy>=2", @@ -24,4 +25,4 @@ tsls = { path = "../../tsls", editable = true } line-length = 89 [tool.ruff.lint.isort] -known-first-party = ["analysis", "thresholding", "tsls", "verification"] +known-first-party = ["analysis", "queries", "thresholding", "tsls", "verification"] diff --git a/experiments/compression/report.py b/experiments/compression/report.py index a605a79..f1cab75 100644 --- a/experiments/compression/report.py +++ b/experiments/compression/report.py @@ -28,12 +28,28 @@ def summarise_matrix(row, num_sites, parsimony_sites): """ ts = tskit.load(RESULTS / row.trees) cm = tsls.load_compressed_matrix(RESULTS / row.matrix, ts) - lists = analysis.list_measures(cm) + num_alleles = getattr(row, "num_alleles", np.nan) + if np.isnan(num_alleles): + # Captured before num_alleles was recorded: sc2ts and msprime use ACGT + num_alleles = 4 + site_map = analysis.SiteMap( + row.algorithm, + rho=row.rho, + mu=row.mu, + n=row.num_nodes if row.weight_by_n else 1, + num_alleles=num_alleles, + ) + lists = analysis.list_measures(cm, site_map) decoded = analysis.decoded_measures( - cm, num_sites=num_sites, parsimony_sites=parsimony_sites + cm, num_sites=num_sites, parsimony_sites=parsimony_sites, site_map=site_map ) nonempty = lists[lists.stored > 0] + # Encoding measures compare each site with the previous non-empty one + carried = nonempty.iloc[1:] + churn = carried.entered + carried.left + first = nonempty.stored.iloc[0] summary = { + "case": row.case, "arg": row.arg, "label": row.label, "query": row.query, @@ -48,11 +64,29 @@ def summarise_matrix(row, num_sites, parsimony_sites): "distinct": nonempty.distinct.mean(), "median_clade_size": decoded.median_clade_size.mean(), "mean_clade_size": decoded.mean_clade_size.mean(), - "samples_only": decoded.samples_only.mean(), + "samples_only": decoded.samples_only.mean() + if "samples_only" in decoded + else np.nan, "jaccard": nonempty.jaccard.mean(), "delta_per_site": (nonempty.added + nonempty.removed).mean(), "mutation_nodes": nonempty.mutation_nodes.mean(), + "churn": churn.mean(), + "changed": carried.changed.mean(), + "changed_tol": carried.changed_tol.mean(), + "affine_residual": carried.affine_residual.mean(), + "affine_fit": carried.affine_fit.mean(), + "affine_work": (churn + carried.affine_residual).mean(), + "total_stored": nonempty.stored.sum(), + "delta_values": first + (churn + carried.changed).sum(), + "affine_delta_values": first + + (churn + carried.affine_residual).sum() + + 2 * len(nonempty), + # The same, on the minimal (change point) lists at the decoded sites + "cp_churn": (decoded.cp_entered + decoded.cp_left).mean(), + "cp_changed": decoded.cp_changed.mean(), + "cp_affine_residual": decoded.cp_affine_residual.mean(), } + summary["cp_affine_work"] = summary["cp_churn"] + summary["cp_affine_residual"] for digits in [1, 2, 4]: summary[f"quantised_{digits}"] = decoded[f"quantised_{digits}"].mean() if row.algorithm == "viterbi": @@ -64,12 +98,24 @@ def summarise_matrix(row, num_sites, parsimony_sites): summary["delta_entries"] = (nonempty.added + nonempty.removed).sum() + ( nonempty.stored.iloc[0] ) + # The traceback stores nodes and flags, not values + summary["flag_changed"] = carried.flag_changed.mean() + summary["delta_traceback"] = first + (churn + carried.flag_changed).sum() return summary, lists, decoded, analysis.value_histogram(cm, top=6) def plot_along_genome(lists, decoded, title, path): fig, ax = plt.subplots(figsize=(10, 3.5)) - ax.plot(lists.site, lists.stored, label="stored", linewidth=0.8) + nonempty = lists[lists.stored > 0] + ax.plot(nonempty.site, nonempty.stored, label="stored", linewidth=0.8) + if "affine_residual" in nonempty: + work = nonempty.entered + nonempty.left + nonempty.affine_residual + ax.plot( + nonempty.site, + work.rolling(25, min_periods=1).mean(), + label="entered + left + affine residual (rolling mean)", + linewidth=0.8, + ) ax.plot( decoded.site, decoded.change_points, @@ -112,6 +158,84 @@ def plot_encodings(summary, path): plt.close(fig) +def plot_work(summary, path): + """ + Plot the per-site work proxy of the current engine (the stored list) + against an affine engine (entries entering or leaving the list, and + values not following the global map), for each matrix. + """ + fig, axes = plt.subplots(1, 2, figsize=(14, 4.5)) + for ax, algorithm in zip(axes, ["viterbi", "forward"]): + d = summary[summary.algorithm == algorithm] + labels = [f"{r.arg}\n{r.label}" for r in d.itertuples()] + x = np.arange(len(d)) + width = 0.35 + ax.bar(x - width / 2, d.stored, width, label="stored (current engine)") + ax.bar(x + width / 2, d.affine_work, width, label="churn + affine residual") + ax.set_yscale("log") + ax.set_xticks(x, labels, rotation=90, fontsize=5) + ax.set_ylabel("Entries per site") + ax.set_title(algorithm) + ax.legend(fontsize=8) + fig.tight_layout() + fig.savefig(path, dpi=100) + plt.close(fig) + + +def encodings_table(summary, algorithm): + """ + Return the markdown table of per-site encoding measures and encoding + sizes for the specified algorithm, with one row per ARG (means over its + queries). + """ + d = summary[summary.algorithm == algorithm].copy() + d["values_delta_saving"] = d.total_stored / d.delta_values + d["values_affine_saving"] = d.total_stored / d.affine_delta_values + d["work_saving"] = d.stored / d.affine_work + columns = [ + "stored", + "churn", + "changed", + "changed_tol", + "affine_residual", + "affine_fit", + "values_delta_saving", + "values_affine_saving", + "work_saving", + ] + if algorithm == "viterbi": + d["traceback_delta_saving"] = d.traceback_entries / d.delta_traceback + columns += ["flag_changed", "traceback_delta_saving"] + return grouped_table(d, columns) + + +def minimal_table(summary, algorithm): + """ + Return the markdown table of encoding measures on the minimal (change + point) lists, at the decoded sites. + """ + d = summary[summary.algorithm == algorithm].copy() + d["cp_delta_work_saving"] = d.change_points / (d.cp_churn + d.cp_changed) + d["cp_affine_work_saving"] = d.change_points / d.cp_affine_work + columns = [ + "stored", + "change_points", + "cp_churn", + "cp_changed", + "cp_affine_residual", + "cp_delta_work_saving", + "cp_affine_work_saving", + ] + return grouped_table(d, columns) + + +def grouped_table(d, columns): + table = d.groupby(["case", "arg"], sort=False)[columns].mean().reset_index() + counts = d.groupby(["case", "arg"], sort=False).size().to_numpy() + table.insert(2, "queries", counts) + return table.to_markdown(index=False, floatfmt=".3g") + + def main(): parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) parser.add_argument( @@ -123,8 +247,17 @@ def main(): default=10, help="Sites for the samples-only parsimony", ) + parser.add_argument( + "--cases", + default=None, + help="Only analyse these cases (comma separated), for quick iteration", + ) args = parser.parse_args() manifest = pd.read_csv(RESULTS / "manifest.csv") + if "case" not in manifest: + manifest["case"] = manifest.arg.str.split("_").str[0] + if args.cases is not None: + manifest = manifest[manifest.case.isin(args.cases.split(","))] rows = [] plots = RESULTS / "plots" plots.mkdir(exist_ok=True) @@ -150,6 +283,7 @@ def main(): summary = pd.DataFrame(rows) summary.to_csv(RESULTS / "summary.csv", index=False) plot_encodings(summary, plots / "encodings.png") + plot_work(summary, plots / "affine_work.png") viterbi = summary[summary.algorithm == "viterbi"].copy() viterbi["delta_saving"] = viterbi.traceback_entries / viterbi.delta_entries @@ -216,6 +350,53 @@ def main(): "", "`quantised_d`: change points after rounding log10 values to 1/d.", "", + "## Encodings", + "", + "Per-site means, comparing each site's stored list with the previous", + "site's, then averaged over the queries of each ARG. `churn` is the", + "number of nodes entering or leaving the list. `changed` counts nodes", + "in both lists whose value changed (`changed_tol`: by a relative", + "difference of more than 1e-9). `affine_residual` counts those whose", + "value is not predicted by the site's global map (Viterbi:", + "max(v stay, rho/n) e / c; forward: a v + b, with b fitted), and", + "`affine_fit` is the fraction that are.", + "", + "Savings are ratios of entries stored over all sites. For values: the", + "full lists against a delta encoding (first list, then churn and", + "changed values), and against an affine + delta encoding (churn and", + "affine residuals, plus two map parameters per site). `work_saving`", + "is the list size against churn + affine residual per site: the", + "values an affine engine would update explicitly. For the Viterbi", + "traceback, which stores nodes and flags but not values, the engine's", + "lists against a delta encoding of nodes and flags.", + "", + "### Viterbi", + "", + encodings_table(summary, "viterbi"), + "", + "### Forward", + "", + encodings_table(summary, "forward"), + "", + "### On the minimal lists", + "", + "The engine stores its lists before removing values equal to the", + "parent's, and gives nodes their own values when edges change, so", + "the stored lists churn more than the values do. These are the same", + "measures on the change points (the minimal lists), comparing each", + "decoded site with the previous one: the entries an engine that kept", + "only change points would update with a delta encoding (churn +", + "changed) or an affine one (churn + affine residual), and the", + "corresponding savings against the change points.", + "", + "#### Viterbi", + "", + minimal_table(summary, "viterbi"), + "", + "#### Forward", + "", + minimal_table(summary, "forward"), + "", ] (RESULTS / "report.md").write_text("\n".join(lines)) print(f"Wrote {RESULTS / 'report.md'}") diff --git a/experiments/compression/uv.lock b/experiments/compression/uv.lock index c4a8cef..64d42b4 100644 --- a/experiments/compression/uv.lock +++ b/experiments/compression/uv.lock @@ -22,6 +22,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, ] +[[package]] +name = "click" +version = "8.5.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c7/0e/7fa0ef50764b67090eca4114772a2abf8b6148198475e54c660b97caeee6/click-8.5.0.tar.gz", hash = "sha256:ba0d2089de75ea0310e2dde03160e6ca10009947fb95a182f9b54021bb272e34", size = 382235, upload-time = "2026-08-26T13:33:14.56Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/58/50/6c0d534c5f134586a8e1ba4e330569e32f057e33372ae556463212fb4cd3/click-8.5.0-py3-none-any.whl", hash = "sha256:255bc9599cf7748b4b1a446ccc735421bd08a2ae529a8b88597d3de5664ee360", size = 125251, upload-time = "2026-08-26T13:33:12.928Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -1156,13 +1165,17 @@ requires-dist = [ ] [package.metadata.requires-dev] -test = [{ name = "pytest" }] +test = [ + { name = "pytest" }, + { name = "pytest-xdist", specifier = ">=3.8.0" }, +] [[package]] name = "tsls-compression-experiments" version = "0.0.0" source = { virtual = "." } dependencies = [ + { name = "click" }, { name = "matplotlib" }, { name = "msprime" }, { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, @@ -1176,6 +1189,7 @@ dependencies = [ [package.metadata] requires-dist = [ + { name = "click", specifier = ">=8.5.0" }, { name = "matplotlib" }, { name = "msprime" }, { name = "numpy", specifier = ">=2" }, diff --git a/tsls/lib/matcher.c b/tsls/lib/matcher.c index 1eb9307..de5a0f1 100644 --- a/tsls/lib/matcher.c +++ b/tsls/lib/matcher.c @@ -561,6 +561,8 @@ ancestor_matcher_alloc(ancestor_matcher_t *self, self->recombination_required = malloc((N + 1) * sizeof(*self->recombination_required)); self->likelihood_cache = malloc((N + 1) * sizeof(*self->likelihood_cache)); + self->recombination_required_cache + = malloc((N + 1) * sizeof(*self->recombination_required_cache)); self->likelihood_nodes = malloc((N + 1) * sizeof(*self->likelihood_nodes)); self->likelihood_nodes_tmp = malloc((N + 1) * sizeof(*self->likelihood_nodes_tmp)); self->values_tmp = malloc((N + 1) * sizeof(*self->values_tmp)); @@ -571,6 +573,7 @@ ancestor_matcher_alloc(ancestor_matcher_t *self, || self->left_sib == NULL || self->right_sib == NULL || self->num_below == NULL || self->num_inheriting == NULL || self->recombination_required == NULL || self->likelihood == NULL || self->likelihood_cache == NULL + || self->recombination_required_cache == NULL || self->likelihood_nodes == NULL || self->likelihood_nodes_tmp == NULL || self->values_tmp == NULL || self->allelic_state == NULL) { ret = TSLS_ERR_NO_MEMORY; @@ -613,6 +616,7 @@ ancestor_matcher_free(ancestor_matcher_t *self) tsi_safe_free(self->recombination_required); tsi_safe_free(self->likelihood); tsi_safe_free(self->likelihood_cache); + tsi_safe_free(self->recombination_required_cache); tsi_safe_free(self->likelihood_nodes); tsi_safe_free(self->likelihood_nodes_tmp); tsi_safe_free(self->values_tmp); @@ -1063,12 +1067,57 @@ ancestor_matcher_insert_mutation_likelihoods( } } +/* As ancestor_matcher_get_inherited_likelihood, but also return the + * recombination_required flag that u inherits, caching it along the path + * with the likelihood. */ +static inline double +ancestor_matcher_get_inherited_likelihood_and_flag( + ancestor_matcher_t *self, tsk_id_t u, int8_t *flag) +{ + const double *restrict L = self->likelihood; + const int8_t *restrict R = self->recombination_required; + double *restrict L_cache = self->likelihood_cache; + int8_t *restrict R_cache = self->recombination_required_cache; + const tsk_id_t *restrict parent = self->parent; + double value; + int8_t r; + tsk_id_t v; + + v = u; + while (likely(L[v] == NULL_LIKELIHOOD) && likely(L_cache[v] == CACHE_UNSET)) { + v = parent[v]; + assert(v != NULL_NODE); + } + if (L_cache[v] == CACHE_UNSET) { + value = L[v]; + r = R[v]; + } else { + value = L_cache[v]; + r = R_cache[v]; + } + assert(value >= 0); + v = u; + while (likely(L[v] == NULL_LIKELIHOOD) && likely(L_cache[v] == CACHE_UNSET)) { + L_cache[v] = value; + R_cache[v] = r; + v = parent[v]; + } + *flag = r; + return value; +} + /* Remove likelihoods that are equal to the likelihood inherited from the - * parent. */ + * parent. If check_recombination_required is true, also require the + * recombination_required flag to equal the one inherited from the parent, + * so that the traceback, in which nodes inherit the flag of their nearest + * ancestor in the list, is unchanged. */ static void -ancestor_matcher_coalesce_likelihoods(ancestor_matcher_t *self) +ancestor_matcher_coalesce_likelihoods( + ancestor_matcher_t *self, bool check_recombination_required) { + const int8_t *restrict R = self->recombination_required; double L_p; + int8_t R_p; tsk_id_t u, p; double *restrict L = self->likelihood; const tsk_id_t *restrict parent = self->parent; @@ -1085,9 +1134,19 @@ ancestor_matcher_coalesce_likelihoods(ancestor_matcher_t *self) if (p != NULL_NODE) { cached_paths[num_cached_paths] = p; num_cached_paths++; - L_p = ancestor_matcher_get_inherited_likelihood(self, p); - if (L[u] == L_p) { - L[u] = NULL_LIKELIHOOD; + if (check_recombination_required) { + /* A removed node had the inherited value and flag, so the + * cached values below it stay valid. */ + L_p = ancestor_matcher_get_inherited_likelihood_and_flag( + self, p, &R_p); + if (L[u] == L_p && R[u] == R_p) { + L[u] = NULL_LIKELIHOOD; + } + } else { + L_p = ancestor_matcher_get_inherited_likelihood(self, p); + if (L[u] == L_p) { + L[u] = NULL_LIKELIHOOD; + } } } if (L[u] >= 0) { @@ -1305,6 +1364,10 @@ ancestor_matcher_process_site_viterbi(ancestor_matcher_t *self, const tsk_id_t s } ancestor_matcher_prune(self); ancestor_matcher_count_distinct(self); + /* Remove redundant nodes before storing the traceback, keeping those + * whose recombination_required flag differs from the inherited one. + * Those are left for the next site, where their flags are recomputed. */ + ancestor_matcher_coalesce_likelihoods(self, true); ret = ancestor_matcher_store_traceback(self, site); if (ret != 0) { goto out; @@ -1316,7 +1379,6 @@ ancestor_matcher_process_site_viterbi(ancestor_matcher_t *self, const tsk_id_t s goto out; } } - ancestor_matcher_coalesce_likelihoods(self); out: return ret; } @@ -1374,7 +1436,7 @@ ancestor_matcher_process_site_forward(ancestor_matcher_t *self, const tsk_id_t s if (ret != 0) { goto out; } - ancestor_matcher_coalesce_likelihoods(self); + ancestor_matcher_coalesce_likelihoods(self, false); out: return ret; } @@ -1436,7 +1498,7 @@ ancestor_matcher_process_site_backward(ancestor_matcher_t *self, const tsk_id_t transition->total = ancestor_matcher_total_likelihood(self); transition->rho = self->recombination_rate[site]; transition->normalisation_factor = forward_normalisation_factor[site]; - ancestor_matcher_coalesce_likelihoods(self); + ancestor_matcher_coalesce_likelihoods(self, false); out: return ret; } diff --git a/tsls/lib/tsls.h b/tsls/lib/tsls.h index f823696..506a238 100644 --- a/tsls/lib/tsls.h +++ b/tsls/lib/tsls.h @@ -147,6 +147,9 @@ typedef struct { int64_t *num_inheriting; double *likelihood; double *likelihood_cache; + /* The recombination_required flag inherited by each node, valid where + * likelihood_cache is set. Only filled when coalescing Viterbi. */ + int8_t *recombination_required_cache; allele_t *allelic_state; int num_likelihood_nodes; /* At each site, record a node with the maximum likelihood. */ From bafe726b333b0be6bf83aa6714070ebd40264d29 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 22 Oct 2026 21:30:38 +0100 Subject: [PATCH 2/2] Speed up Viterbi allele lookup and first-tree build in tsls tsls stores a smaller traceback than tsinfer-hmm but was still slower, because storing the traceback is under 1% of the time in both engines. Profiling both on the same 1000 Genomes chr20p queries (callgrind, limited to ancestor_matcher_find_path) showed tsls doing 1.26x (samples) to 1.49x (late ancestors) the instructions, mostly in the allele lookup and in building the first tree of each query. What was changed, and how it landed: 1. Allele lookup. tsinfer-hmm sets root 0 to the ancestral state, so its walk up the tree checks only for a set allele; tsls allows several roots and checked for the root at every step. tsls now sets every root's allele (roots are all in the likelihood list) before each site's update, and the Viterbi, forward and backward loops use get_marked_allelic_state. A first version kept an assert in the walk; because the Python module is built with asserts on, it cost as much as the check it replaced (12.71 G against 12.50 G instructions on one sample), so the walk has no assert. The traceback keeps get_allelic_state. 2. First-tree build. For Viterbi, init_sweep now builds the first tree with plain edge insertions and removals (ancestor_matcher_next_tree_topology), as tsinfer-hmm does, and no longer initialises num_below. Forward and backward, which track num_below, are unchanged. Instructions, tsinfer-hmm / tsls before / tsls after: 20 group 1485 ancestors 6.65 / 9.89 / 7.44 G; one sample 9.92 / 12.50 / 10.96 G. Paired wall times (best of 3, alternating engines per query), tsls over tsinfer-hmm: group 1485 ancestors 1.35 -> 1.10, samples 1.06 -> 1.06. Change 2 is the win; change 1 saves instructions but no time, because the walk is limited by memory access. run.py ratios proved unreliable for comparing engines, since it runs the engines in blocks and machine speed drifts; the README documents this and the paired method. Paths are unchanged: tsls still finds tsinfer's path on every query checked, and the tsls, C and benchmarking tests pass. --- benchmarking/README.md | 76 +++++++++++++++++++- tsls/lib/matcher.c | 159 +++++++++++++++++++++++++++++++++++++---- 2 files changed, 221 insertions(+), 14 deletions(-) diff --git a/benchmarking/README.md b/benchmarking/README.md index 74abc70..7d7dd68 100644 --- a/benchmarking/README.md +++ b/benchmarking/README.md @@ -363,10 +363,84 @@ for the whole group divided by the number of haplotypes. not yet in the tree sequence. - tsls is 3–10% slower than tsinfer-hmm in mean time per query, with the same matcher memory (64 MiB for ancestors, about 90 MiB for samples) - and traceback sizes. Building tsls's `MatcherIndexes` takes several + and traceback sizes. (These figures predate the engine changes + described below, and `run.py` timings are biased by run order; see + [tsls against tsinfer-hmm](#tsls-against-tsinfer-hmm-where-the-time-goes).) + Building tsls's `MatcherIndexes` takes several times longer (3.6 s against 0.4 s for the full ancestors tree sequence), but this is done once per group. - Per-query times are close to tsinfer's per-haplotype CPU time, so the HMM accounts for most of tsinfer's matching cost. - The cost is in the last few ancestor groups and the samples. Matching all 6,404 samples takes about 4.5 CPU hours with either engine. + +### tsls against tsinfer-hmm: where the time goes + +tsls now stores a smaller Viterbi traceback than tsinfer-hmm (13 against +21 entries per site for samples; see `experiments/compression/README.md`), +but it isn't faster, because storing the traceback is under 1% of the +time in both engines. Both engines update the same likelihood lists at +each site, since tsinfer-hmm coalesces them right after storing. The time +goes into walking up the tree, to find each listed node's allele and the +value it inherits, and into moving from tree to tree. Building the first +tree of a query means processing most of the edges before its window, +which matters for late ancestors. + +Instruction counts (callgrind, collection limited to +`ancestor_matcher_find_path`, 20 group 1485 ancestors and 1 sample) showed +where tsls did more work: + +| | tsinfer-hmm | tsls before | tsls after | +|---|---|---|---| +| 20 ancestors, group 1485 | 6.65 G | 9.89 G | 7.44 G | +| 1 sample | 9.92 G | 12.50 G | 10.96 G | + +Two changes to `tsls/lib/matcher.c` account for the difference between +"before" and "after": + +1. **Allele lookup.** tsinfer-hmm's trees have the single root 0, whose + allele it sets to the ancestral state, so the walk up the tree only + checks for a set allele. tsls allows several roots, so + `get_allelic_state` also checked for the root at each step. tsls now + sets every root's allele to the ancestral state before the site's + update (the roots all have likelihoods, so they are found in the + likelihood list), and the three likelihood loops use + `get_marked_allelic_state`, which walks without the check. The walk + deliberately has no `assert`: the Python module is compiled with + asserts on, and an `assert` there cost as much as the check it + replaced (12.71 G on the sample, against 12.50 G before). The traceback + keeps the general `get_allelic_state`, since its tree has no + likelihoods. +2. **Building the first tree.** For Viterbi, `init_sweep` now builds the + first tree with plain edge insertions and removals + (`ancestor_matcher_next_tree_topology`), as tsinfer-hmm does, instead + of passing every edge through the general wrappers that handle + likelihoods, subtree counts and isolated nodes. It also no longer + initialises `num_below`, which only forward and backward use. Forward + and backward still build it the general way, because they track + `num_below`. + +Wall times, measured in pairs: each query is run three times with each +engine, alternating, and the best time is kept. Ratios are total tsls time +over total tsinfer-hmm time (100 group 1485 ancestors, 6 samples): + +| | tsls before ÷ tsinfer-hmm | tsls after ÷ tsinfer-hmm | +|---|---|---| +| ancestors, group 1485 | 1.35 | 1.10 | +| samples | 1.06 | 1.06 | + +- Change 2 is a clear win on late ancestors. +- Change 1 removes about 1.5 G instructions per sample but no wall time. + The walk up the tree is limited by memory access (each step loads + `allelic_state` and `parent` for a node that is usually not in cache), + so the extra compare was nearly free. +- What is left is mostly the same kind of overhead in the coalescing and + cache code: `ancestor_matcher_get_inherited_likelihood` and + `ancestor_matcher_reset_likelihood_cache` have asserts and root checks + inside their walks, which tsinfer-hmm's equivalents don't. By the + result for change 1, these would also save instructions more than time. +- Timing caveat: `run.py` runs all queries for one engine, then all for + the next, so drift in machine speed (large on this VM: the same query + can vary by 20% between passes) shows up as a difference between + engines. Ratios from `run.py` ranged from 1.09 to 1.25 on the same + builds. Use paired measurements, as above, to compare engines. diff --git a/tsls/lib/matcher.c b/tsls/lib/matcher.c index de5a0f1..92f1105 100644 --- a/tsls/lib/matcher.c +++ b/tsls/lib/matcher.c @@ -99,6 +99,22 @@ get_allelic_state(const allele_t *restrict allelic_state, return allelic_state[u]; } +/* As get_allelic_state, but for a tree in which every root has its + * allelic state set (see ancestor_matcher_set_root_allelic_state), so the + * walk up the tree doesn't need to check for the root. This is the + * innermost loop of the matcher, so there is deliberately no assert here + * (the Python module is compiled with asserts on); the roots are checked + * in ancestor_matcher_check_state. */ +static inline allele_t +get_marked_allelic_state( + const allele_t *restrict allelic_state, const tsk_id_t *restrict parent, tsk_id_t u) +{ + while (allelic_state[u] == TSK_NULL) { + u = parent[u]; + } + return allelic_state[u]; +} + static inline bool is_integer(double x) { @@ -984,6 +1000,74 @@ ancestor_matcher_next_tree( } } +/* Move to the next tree in the direction of travel, updating only the tree + * topology. This is the same as ancestor_matcher_next_tree with + * update_likelihoods false when num_below isn't tracked, but avoids its + * per-edge bookkeeping. It is used to build the first tree of a sweep, + * which may mean processing most of the edges. */ +static void +ancestor_matcher_next_tree_topology(ancestor_matcher_t *self, tree_position_t *tp) +{ + const matcher_indexes_t *mi = self->matcher_indexes; + const int_fast32_t M = (int_fast32_t) mi->num_edges; + const coordinate_t sequence_length = mi->sites.position[mi->num_sites]; + const edge_t *restrict left_index = mi->left_index_edges; + const edge_t *restrict right_index = mi->right_index_edges; + tsk_id_t *restrict parent = self->parent; + tsk_id_t *restrict left_child = self->left_child; + tsk_id_t *restrict right_child = self->right_child; + tsk_id_t *restrict left_sib = self->left_sib; + tsk_id_t *restrict right_sib = self->right_sib; + int_fast32_t in_index = tp->in_index; + int_fast32_t out_index = tp->out_index; + coordinate_t pos; + + assert(!self->track_num_below); + if (tp->direction > 0) { + pos = tp->right; + while (out_index < M && right_index[out_index].right == pos) { + remove_edge(right_index[out_index], parent, left_child, right_child, + left_sib, right_sib); + out_index++; + } + while (in_index < M && left_index[in_index].left == pos) { + insert_edge(left_index[in_index], parent, left_child, right_child, left_sib, + right_sib); + in_index++; + } + tp->left = pos; + tp->right = sequence_length; + if (in_index < M) { + tp->right = TSK_MIN(tp->right, left_index[in_index].left); + } + if (out_index < M) { + tp->right = TSK_MIN(tp->right, right_index[out_index].right); + } + } else { + pos = tp->left; + while (out_index >= 0 && left_index[out_index].left == pos) { + remove_edge(left_index[out_index], parent, left_child, right_child, + left_sib, right_sib); + out_index--; + } + while (in_index >= 0 && right_index[in_index].right == pos) { + insert_edge(right_index[in_index], parent, left_child, right_child, + left_sib, right_sib); + in_index--; + } + tp->right = pos; + tp->left = 0; + if (out_index >= 0) { + tp->left = TSK_MAX(tp->left, left_index[out_index].left); + } + if (in_index >= 0) { + tp->left = TSK_MAX(tp->left, right_index[in_index].right); + } + } + tp->in_index = in_index; + tp->out_index = out_index; +} + /* Reset the matcher and build the tree containing the specified position, * giving each root the specified initial likelihood. */ static void @@ -1001,13 +1085,21 @@ ancestor_matcher_init_sweep(ancestor_matcher_t *self, tree_position_t *tp, memset(self->right_sib, 0xff, N * sizeof(*self->right_sib)); memset(self->recombination_required, 0xff, N * sizeof(*self->recombination_required)); for (j = 0; j < N; j++) { - self->num_below[j] = 1; self->likelihood_cache[j] = CACHE_UNSET; } tree_position_init(tp, self->matcher_indexes, direction); - do { - ancestor_matcher_next_tree(self, tp, false); - } while (!(tp->left <= position && position < tp->right)); + if (self->track_num_below) { + for (j = 0; j < N; j++) { + self->num_below[j] = 1; + } + do { + ancestor_matcher_next_tree(self, tp, false); + } while (!(tp->left <= position && position < tp->right)); + } else { + do { + ancestor_matcher_next_tree_topology(self, tp); + } while (!(tp->left <= position && position < tp->right)); + } self->num_likelihood_nodes = 0; for (u = 0; u < (tsk_id_t) N; u++) { @@ -1043,6 +1135,47 @@ ancestor_matcher_unset_allelic_state(ancestor_matcher_t *self, const tsk_id_t si } } +/* Set the allelic state of every root to the ancestral state, so that + * get_marked_allelic_state finds a state for every node in the tree. + * Every root has a likelihood, so the roots are found in the likelihood + * list. This must be called before ancestor_matcher_set_allelic_state, so + * that mutations above a root take precedence, and the likelihood list + * must not change before ancestor_matcher_unset_root_allelic_state. */ +static inline void +ancestor_matcher_set_root_allelic_state(ancestor_matcher_t *self, const tsk_id_t site) +{ + const allele_t ancestral_state = self->matcher_indexes->sites.ancestral_state[site]; + const tsk_id_t *restrict L_nodes = self->likelihood_nodes; + const tsk_id_t *restrict parent = self->parent; + allele_t *restrict allelic_state = self->allelic_state; + tsk_id_t u; + int j; + + for (j = 0; j < self->num_likelihood_nodes; j++) { + u = L_nodes[j]; + if (parent[u] == NULL_NODE) { + allelic_state[u] = ancestral_state; + } + } +} + +static inline void +ancestor_matcher_unset_root_allelic_state(ancestor_matcher_t *self) +{ + const tsk_id_t *restrict L_nodes = self->likelihood_nodes; + const tsk_id_t *restrict parent = self->parent; + allele_t *restrict allelic_state = self->allelic_state; + tsk_id_t u; + int j; + + for (j = 0; j < self->num_likelihood_nodes; j++) { + u = L_nodes[j]; + if (parent[u] == NULL_NODE) { + allelic_state[u] = TSK_NULL; + } + } +} + /* Make sure that every node with a mutation at this site that is in the * tree has its own likelihood. */ static void @@ -1308,7 +1441,6 @@ ancestor_matcher_process_site_viterbi(ancestor_matcher_t *self, const tsk_id_t s const tsk_id_t *restrict L_nodes = self->likelihood_nodes; const tsk_id_t *restrict parent = self->parent; const allele_t *restrict allelic_state = self->allelic_state; - const allele_t ancestral_state = self->matcher_indexes->sites.ancestral_state[site]; double *restrict L = self->likelihood; int8_t *restrict recombination_required = self->recombination_required; const double rho = self->recombination_rate[site]; @@ -1320,6 +1452,7 @@ ancestor_matcher_process_site_viterbi(ancestor_matcher_t *self, const tsk_id_t s int j; ancestor_matcher_insert_mutation_likelihoods(self, site); + ancestor_matcher_set_root_allelic_state(self, site); ancestor_matcher_set_allelic_state(self, site); max_L = -1; max_L_node = NULL_NODE; @@ -1335,8 +1468,7 @@ ancestor_matcher_process_site_viterbi(ancestor_matcher_t *self, const tsk_id_t s recombination_required[u] = true; } p_e = emission_probability( - get_allelic_state(allelic_state, parent, u, ancestral_state), state, mu, - num_alleles); + get_marked_allelic_state(allelic_state, parent, u), state, mu, num_alleles); L[u] = p_t * p_e; if (L[u] > max_L) { max_L = L[u]; @@ -1344,6 +1476,7 @@ ancestor_matcher_process_site_viterbi(ancestor_matcher_t *self, const tsk_id_t s } } ancestor_matcher_unset_allelic_state(self, site); + ancestor_matcher_unset_root_allelic_state(self); (void) num_likelihood_nodes; if (max_L <= 0) { @@ -1395,7 +1528,6 @@ ancestor_matcher_process_site_forward(ancestor_matcher_t *self, const tsk_id_t s const tsk_id_t *restrict L_nodes = self->likelihood_nodes; const tsk_id_t *restrict parent = self->parent; const allele_t *restrict allelic_state = self->allelic_state; - const allele_t ancestral_state = self->matcher_indexes->sites.ancestral_state[site]; double *restrict L = self->likelihood; const double rho = self->recombination_rate[site]; const double mu = self->mismatch_rate[site]; @@ -1411,15 +1543,16 @@ ancestor_matcher_process_site_forward(ancestor_matcher_t *self, const tsk_id_t s if (total < 0) { total = ancestor_matcher_total_likelihood(self); } + ancestor_matcher_set_root_allelic_state(self, site); ancestor_matcher_set_allelic_state(self, site); for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; p_e = emission_probability( - get_allelic_state(allelic_state, parent, u, ancestral_state), state, mu, - num_alleles); + get_marked_allelic_state(allelic_state, parent, u), state, mu, num_alleles); L[u] = (L[u] * (1 - rho) + rho / n * total) * p_e; } ancestor_matcher_unset_allelic_state(self, site); + ancestor_matcher_unset_root_allelic_state(self); normalisation_factor = ancestor_matcher_total_likelihood(self); if (!(normalisation_factor > 0)) { ret = TSLS_ERR_MATCH_IMPOSSIBLE; @@ -1461,7 +1594,6 @@ ancestor_matcher_process_site_backward(ancestor_matcher_t *self, const tsk_id_t const tsk_id_t *restrict L_nodes = self->likelihood_nodes; const tsk_id_t *restrict parent = self->parent; const allele_t *restrict allelic_state = self->allelic_state; - const allele_t ancestral_state = self->matcher_indexes->sites.ancestral_state[site]; double *restrict L = self->likelihood; const double mu = self->mismatch_rate[site]; const double n = ancestor_matcher_get_n(self); @@ -1486,14 +1618,15 @@ ancestor_matcher_process_site_backward(ancestor_matcher_t *self, const tsk_id_t goto out; } ancestor_matcher_compute_num_inheriting(self); + ancestor_matcher_set_root_allelic_state(self, site); ancestor_matcher_set_allelic_state(self, site); for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; L[u] *= emission_probability( - get_allelic_state(allelic_state, parent, u, ancestral_state), state, mu, - num_alleles); + get_marked_allelic_state(allelic_state, parent, u), state, mu, num_alleles); } ancestor_matcher_unset_allelic_state(self, site); + ancestor_matcher_unset_root_allelic_state(self); transition->pending = true; transition->total = ancestor_matcher_total_likelihood(self); transition->rho = self->recombination_rate[site];