From 4d2a67b15ae58440b610332ba14464b4a158f762 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 22 Oct 2026 20:44:48 +0100 Subject: [PATCH 01/25] 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 02/25] 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]; From 7277b871388e19b12cf4f78ef555a02d586c19f0 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 11:35:28 +0100 Subject: [PATCH 03/25] Add a callgrind harness for instruction counts on the tsinfer workloads benchmarking/callgrind.py runs the first queries of each tsinfer chr20p workload with one engine under callgrind, collecting only inside the engine's find_path call, and prints total Ir per workload and engine. Baseline before the tsls_hmm_t refactor (tsls at bafe726), 20 ancestors per group and 1 sample, paths identical between engines: | workload | queries | tsinfer | tsls | ratio | |---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.608 G | 2.055 | | ancestors_g1000 | 20 | 0.621 G | 0.975 G | 1.571 | | ancestors_g1483 | 20 | 2.388 G | 2.845 G | 1.192 | | ancestors_g1484 | 20 | 9.911 G | 11.137 G | 1.124 | | ancestors_g1485 | 20 | 6.653 G | 7.445 G | 1.119 | | samples | 1 | 9.922 G | 10.955 G | 1.104 | On g84, 57% of tsls's instructions are memset: resetting per-site and per-node arrays on every query, which dominates for short ancestors. --- benchmarking/callgrind.py | 166 ++++++++++++++++++++++++++++++++++ tsls/lib/{matcher.c => hmm.c} | 0 2 files changed, 166 insertions(+) create mode 100644 benchmarking/callgrind.py rename tsls/lib/{matcher.c => hmm.c} (100%) diff --git a/benchmarking/callgrind.py b/benchmarking/callgrind.py new file mode 100644 index 0000000..9d17c1b --- /dev/null +++ b/benchmarking/callgrind.py @@ -0,0 +1,166 @@ +""" +Count the instructions the HMM engines execute on the tsinfer workloads. + +``run`` matches the first queries of a workload with one engine. ``measure`` +runs ``run`` under callgrind, with collection limited to the engine's +find_path call, so that the counts are deterministic and leave out setup +(loading the data and building the engine). It prints a markdown table of +total instructions (Ir) for each workload and engine. +""" + +import concurrent.futures +import pathlib +import subprocess +import sys + +import click +import numpy as np + +import engines +import queries +import run as run_module + +# The C function that runs one Viterbi match in each engine. Collection is +# switched on only inside these. +FIND_PATH_FUNCTIONS = { + "tsinfer": ["ancestor_matcher_find_path"], + "tsls": ["tsls_hmm_find_path", "ancestor_matcher_find_path"], +} + +DEFAULT_WORKLOADS = [ + ("ancestors_g84", 20), + ("ancestors_g1000", 20), + ("ancestors_g1483", 20), + ("ancestors_g1484", 20), + ("ancestors_g1485", 20), + ("samples", 1), +] + + +def make_engine(name, ts, model): + return engines.ENGINES[name]( + ts, + mu=model["mu"], + rho=model["rho"], + weight_by_n=model["weight_by_n"], + alleles=model["alleles"], + ) + + +@click.group() +def cli(): + pass + + +@cli.command() +@click.argument("workload_dir", type=click.Path(exists=True, file_okay=False)) +@click.argument("engine_name", type=click.Choice(list(FIND_PATH_FUNCTIONS))) +@click.option("-n", "--num-queries", type=int, default=20) +def run(workload_dir, engine_name, num_queries): + """ + Match the first NUM_QUERIES queries in WORKLOAD_DIR with the specified + engine, printing the parents of each path. + """ + workload_dir = pathlib.Path(workload_dir) + model = queries.load_model(workload_dir) + ts = run_module.load_ts(workload_dir / "reference.trees") + H, _, df = queries.load_queries(workload_dir) + engine = make_engine(engine_name, ts, model) + for query_id in df.query_id[:num_queries]: + start, end = queries.query_window(H, query_id) + _, _, parent, _, _ = engine.match(H[query_id], start, end) + print(query_id, " ".join(str(p) for p in np.asarray(parent))) + + +def read_total_ir(path): + with open(path) as f: + for line in f: + if line.startswith("totals:") or line.startswith("summary:"): + return int(line.split()[1]) + raise ValueError(f"No totals in {path}") + + +def callgrind(workload_dir, engine_name, num_queries, out_dir): + workload = pathlib.Path(workload_dir).name + out_file = pathlib.Path(out_dir) / f"callgrind.{workload}.{engine_name}.out" + toggles = [f"--toggle-collect={f}" for f in FIND_PATH_FUNCTIONS[engine_name]] + cmd = [ + "valgrind", + "--tool=callgrind", + "--collect-atstart=no", + *toggles, + f"--callgrind-out-file={out_file}", + sys.executable, + __file__, + "run", + str(workload_dir), + engine_name, + "-n", + str(num_queries), + ] + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(f"{' '.join(cmd)} failed:\n{result.stderr[-2000:]}") + return read_total_ir(out_file), result.stdout + + +@cli.command() +@click.option( + "--data-dir", + type=click.Path(exists=True, file_okay=False), + default="data/tgp_chr20p", +) +@click.option("-w", "--workload", "workload_names", multiple=True) +@click.option( + "-e", + "--engine", + "engine_names", + multiple=True, + type=click.Choice(list(FIND_PATH_FUNCTIONS)), +) +@click.option("-o", "--out-dir", type=click.Path(file_okay=False), default="callgrind") +@click.option("-j", "--jobs", type=int, default=3) +def measure(data_dir, workload_names, engine_names, out_dir, jobs): + """ + Count the instructions in find_path for each workload and engine. The + callgrind output files are kept in OUT_DIR for callgrind_annotate. + """ + workloads = DEFAULT_WORKLOADS + if len(workload_names) > 0: + workloads = [w for w in DEFAULT_WORKLOADS if w[0] in workload_names] + if len(engine_names) == 0: + engine_names = list(FIND_PATH_FUNCTIONS) + pathlib.Path(out_dir).mkdir(parents=True, exist_ok=True) + tasks = [ + (workload, num_queries, engine) + for workload, num_queries in workloads + for engine in engine_names + ] + with concurrent.futures.ThreadPoolExecutor(jobs) as executor: + futures = { + task: executor.submit( + callgrind, pathlib.Path(data_dir) / task[0], task[2], task[1], out_dir + ) + for task in tasks + } + results = {task: future.result() for task, future in futures.items()} + + header = "| workload | queries | " + " | ".join(engine_names) + " |" + if len(engine_names) == 2: + header += " ratio |" + click.echo(header) + click.echo("|---" * (header.count("|") - 1) + "|") + for workload, num_queries in workloads: + counts = [results[(workload, num_queries, e)][0] for e in engine_names] + row = f"| {workload} | {num_queries} | " + row += " | ".join(f"{c / 1e9:.3f} G" for c in counts) + " |" + if len(engine_names) == 2: + row += f" {counts[1] / counts[0]:.3f} |" + click.echo(row) + paths = {results[(workload, num_queries, e)][1] for e in engine_names} + if len(paths) > 1: + click.echo(f"WARNING: engines found different paths for {workload}") + + +if __name__ == "__main__": + cli() diff --git a/tsls/lib/matcher.c b/tsls/lib/hmm.c similarity index 100% rename from tsls/lib/matcher.c rename to tsls/lib/hmm.c From 023c1fb20041c0aa6725c2adc2ca7f4d67b3efe3 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:08:37 +0100 Subject: [PATCH 04/25] Refactor the tsls engine into tsls_hmm_t on a tsk_treeseq_t The engine (ancestor_matcher_t in matcher.c) becomes tsls_hmm_t in hmm.c, and works directly on a borrowed tsk_treeseq_t instead of the copied edge, site and mutation data in matcher_indexes_t, which is removed. - Trees are visited with tskit's tsk_tree_position_t. The first tree of a sweep is found from a null position with tsk_tree_position_seek_forward or _seek_backward (whichever end is closer, as tsk_tree_seek_from_null does), inserting only the edges that cover the first site, instead of applying every edge from the start (or end) of the sequence. The traceback does the same from the last site of the window, instead of walking in from the end of the sequence. - Alleles are read from the site and mutation tables and encoded by the HMM: by default numbered at each site as tskit does (ancestral 0, then derived states in order of first appearance), with at least two alleles per site so that a tsinfer ancestor's focal sites (which have no mutations yet) can still be matched; or as indexes into a given list (e.g. ACGT for sc2ts). encode_alleles and the ancestral/derived state arrays are gone, along with the error codes tskit now guarantees. - Coordinates are doubles; integer coordinates are no longer required, and path_left/path_right are doubles. - Kept as simple as possible: the root-marked allele lookup and the topology-only first tree from bafe726 are dropped (the seek replaces the latter). The core algorithm (likelihood lists, coalescing with flags, caches, traceback reuse, pruning, forward/backward) is unchanged. Python: MatcherIndexes is replaced by tsls.TreeSequence (a minimal wrapper of tsk_treeseq_t built from a LightweightTableCollection, shareable between HMMs), and AncestorMatcher by tsls.Hmm(ts, recombination, mismatch, *, alleles=None, ...). The extension now compiles all of tskit's C library. Callers in benchmarking, verification and experiments are updated. Tests: new C and Python tests for allele encoding, windows starting and ending throughout the sequence against the dense optimum (exercising seeks from both ends), non-integer coordinates, and sites without mutations. tsls finds tsinfer's path for all 1,914 queries in the six chr20p workloads; traceback sizes and memory are unchanged. Instruction counts (callgrind, find_path only) on the reduced profiling set (see the next commit): | workload | queries | tsinfer | tsls | ratio | |---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.486 G | 1.644 | | ancestors_g1000 | 20 | 0.621 G | 0.686 G | 1.105 | | ancestors_g1483 | 10 | 1.355 G | 1.381 G | 1.019 | | ancestors_g1484 | 3 | 0.646 G | 0.736 G | 1.139 | | ancestors_g1485 | 5 | 1.131 G | 1.117 G | 0.988 | | samples | 1 (100000 sites) | 2.064 G | 2.578 G | 1.249 | On the full set the old engine's ratios were 2.06, 1.57, 1.19, 1.12, 1.12 and 1.10, so the seek more than pays for the dropped special cases on ancestors, while the sample (where the first tree is cheap anyway) is worse, mostly from dropping the root-marked allele lookup. --- README.md | 18 +- benchmarking/engines.py | 53 +- experiments/compression/README.md | 2 +- experiments/compression/analysis.py | 2 +- experiments/compression/capture.py | 18 +- experiments/thresholding/README.md | 2 +- experiments/thresholding/thresholding.py | 14 +- tsls/_tslsmodule.c | 427 ++++---- tsls/lib/err.c | 20 +- tsls/lib/err.h | 14 +- tsls/lib/hmm.c | 1155 ++++++++-------------- tsls/lib/meson.build | 2 +- tsls/lib/tests/tests.c | 246 ++--- tsls/lib/tsls.h | 89 +- tsls/setup.py | 15 +- tsls/tests/test_matching.py | 147 ++- tsls/tsls/__init__.py | 5 +- tsls/tsls/matching.py | 97 +- verification/verification.py | 11 +- 19 files changed, 1022 insertions(+), 1315 deletions(-) diff --git a/README.md b/README.md index da9ad37..1690df6 100644 --- a/README.md +++ b/README.md @@ -29,17 +29,23 @@ module `_tsls`). It is based on the tsinfer engine: node from one site to the next without a recombination requires it to be in the tree over the whole interval between them. - It supports arbitrary topologies (multiple roots, no special root node). -- It provides the Viterbi algorithm (`AncestorMatcher.find_path`, with - the same interface as tsinfer-hmm) and the forward and backward - algorithms (`forward_matrix`, `backward_matrix`), which return - compressed matrices. +- It works directly on a tskit tree sequence (`tsls.TreeSequence`, which + can be shared between HMMs), and moves between trees with tskit's + `tsk_tree_position_t`, seeking straight to the first tree of a window. +- It provides the Viterbi algorithm (`Hmm.find_path`, with the same + interface as tsinfer-hmm's `AncestorMatcher.find_path`) and the forward + and backward algorithms (`forward_matrix`, `backward_matrix`), which + return compressed matrices. - It uses tsinfer's compression (a node's value is dropped when it equals its parent's), a likelihood threshold and `weight_by_n`. -- It supports multiple mutations per site and any number of alleles. +- It supports multiple mutations per site and any number of alleles, + read from the tree sequence and numbered at each site as tskit does + (or by their index in a given list of alleles), and any (including + non-integer) coordinates. Paths match tsinfer-hmm except for ties. To match tsinfer's convention that its ultimate root node 0 is never copied, remove node 0's edges -before building the `MatcherIndexes`; sc2ts ARGs already have an +before building the `tsls.TreeSequence`; sc2ts ARGs already have an isolated node 0. The `benchmarking/` directory compares the engines on the sc2ts ARG, and `verification/` runs statistical validations of tsls on simulated ARGs (plots, not pass/fail tests). diff --git a/benchmarking/engines.py b/benchmarking/engines.py index 99b3274..48636b4 100644 --- a/benchmarking/engines.py +++ b/benchmarking/engines.py @@ -106,7 +106,8 @@ def path_haplotype(ts, left, right, parent): class AncestorMatcherEngine: """ Base class for the engines with tsinfer's AncestorMatcher interface, - which differ only in how the MatcherIndexes are made. + which differ only in how the shared reference (tsinfer's MatcherIndexes + or a tsls TreeSequence) and the per-query matcher are made. If mismatch_threshold is specified, normalised likelihoods are not allowed to drop below that of mismatch_threshold mismatches, as in @@ -115,12 +116,11 @@ class AncestorMatcherEngine: By default, states are encoded as indexes into ALLELES. If alleles is None, the ancestral state of every site is 0 and every mutation has - derived state 1, as in tsinfer. weight_by_n is False as in sc2ts; - tsinfer uses True. + derived state 1, as in tsinfer (tsls numbers the alleles at each site + as tskit does, which is the same for tsinfer's data). weight_by_n is + False as in sc2ts; tsinfer uses True. """ - module = None - def __init__( self, ts, @@ -140,9 +140,13 @@ def __init__( self.likelihood_threshold = mu**mismatch_threshold * 0.99 self.recombination = np.full(ts.num_sites, rho) self.mismatch = np.full(ts.num_sites, mu) - self.matcher_indexes = self.make_matcher_indexes(ts, alleles) + self.alleles = alleles + self.reference = self.make_reference(ts, alleles) + + def make_reference(self, ts, alleles): + raise NotImplementedError - def make_matcher_indexes(self, ts, alleles): + def make_matcher(self): raise NotImplementedError def match(self, h, start=0, end=None): @@ -152,13 +156,7 @@ def match(self, h, start=0, end=None): """ if end is None: end = self.num_sites - matcher = self.module.AncestorMatcher( - self.matcher_indexes, - self.recombination, - self.mismatch, - likelihood_threshold=self.likelihood_threshold, - weight_by_n=self.weight_by_n, - ) + matcher = self.make_matcher() matched = np.zeros(self.num_sites, dtype=np.int8) left, right, parent = matcher.find_path(h.astype(np.int8), start, end, matched) stats = { @@ -176,11 +174,19 @@ def match(self, h, start=0, end=None): class TsinferEngine(AncestorMatcherEngine): name = "tsinfer" - module = tsinfer_hmm - def make_matcher_indexes(self, ts, alleles): + def make_reference(self, ts, alleles): return tsinfer_hmm.MatcherIndexes(ts, vestigial_root=False, alleles=alleles) + def make_matcher(self): + return tsinfer_hmm.AncestorMatcher( + self.reference, + self.recombination, + self.mismatch, + likelihood_threshold=self.likelihood_threshold, + weight_by_n=self.weight_by_n, + ) + class TskitEngine: name = "tskit" @@ -243,12 +249,21 @@ class TslsEngine(AncestorMatcherEngine): """ name = "tsls" - module = tsls - def make_matcher_indexes(self, ts, alleles): + def make_reference(self, ts, alleles): tables = ts.dump_tables() tables.edges.keep_rows(tables.edges.parent != 0) - return tsls.MatcherIndexes(tables.tree_sequence(), alleles=alleles) + return tsls.TreeSequence(tables.tree_sequence()) + + def make_matcher(self): + return tsls.Hmm( + self.reference, + self.recombination, + self.mismatch, + alleles=self.alleles, + likelihood_threshold=self.likelihood_threshold, + weight_by_n=self.weight_by_n, + ) ENGINES = {engine.name: engine for engine in [TsinferEngine, TskitEngine, TslsEngine]} diff --git a/experiments/compression/README.md b/experiments/compression/README.md index f4d1185..d043324 100644 --- a/experiments/compression/README.md +++ b/experiments/compression/README.md @@ -15,7 +15,7 @@ uv run python report.py # results/report.md `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 +`capture.py` uses `Hmm.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 `tsls.load_compressed_matrix(path, ts)`. diff --git a/experiments/compression/analysis.py b/experiments/compression/analysis.py index db72b5c..a16d017 100644 --- a/experiments/compression/analysis.py +++ b/experiments/compression/analysis.py @@ -1,6 +1,6 @@ """ Measures of likelihood compression for a tsls CompressedMatrix (from -AncestorMatcher.viterbi_matrix or forward_matrix) on its tree sequence. +Hmm.viterbi_matrix or forward_matrix) on its tree sequence. The engine stores, at each site, a list of (node, value) pairs, and every node in the tree takes the value of its nearest ancestor in the list. Because diff --git a/experiments/compression/capture.py b/experiments/compression/capture.py index b54cafa..d809ab6 100644 --- a/experiments/compression/capture.py +++ b/experiments/compression/capture.py @@ -61,10 +61,10 @@ def sc2ts_queries(): 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) + mi = tsls.TreeSequence(ts) 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 + ts, mu, rho, weight_by_n=weight_by_n, count_distinct=True, tree_sequence=mi ) (_, _, parent), matched, vm = am.viterbi_matrix(h) fm = am.forward_matrix(h) @@ -148,8 +148,8 @@ def check_window(mi, ts, h, start, end, vm, rho, mu, path): which validates the windowing used for forward. """ wts = window_ts(ts, start, end) - am = tsls.AncestorMatcher( - tsls.MatcherIndexes(wts), + am = tsls.Hmm( + wts, np.full(wts.num_sites, rho), np.full(wts.num_sites, mu), ) @@ -179,15 +179,13 @@ def capture_tsinfer(workload, *, num_queries, forward_sites, seed, check, rows): ts = drop_root_edges(tskit.load(query_dir / "reference.trees")) ts_path = RESULTS / f"{name}.trees" ts.dump(ts_path) - mi = tsls.MatcherIndexes(ts) + mi = tsls.TreeSequence(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) - ) + am = tsls.Hmm(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)) @@ -235,8 +233,8 @@ def capture_tsinfer(workload, *, num_queries, forward_sites, seed, check, rows): 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), + fam = tsls.Hmm( + wts, np.full(wts.num_sites, rho), np.full(wts.num_sites, mu), ) diff --git a/experiments/thresholding/README.md b/experiments/thresholding/README.md index f55a2db..8573782 100644 --- a/experiments/thresholding/README.md +++ b/experiments/thresholding/README.md @@ -34,7 +34,7 @@ reached again by recombination. changes no Viterbi likelihoods. With `weight_by_n=False`, as sc2ts uses, `t* = ρ`. -`AncestorMatcher(count_distinct=True)` also records the mean number of +`Hmm(count_distinct=True)` also records the mean number of **distinct** values per site in Viterbi (`mean_num_distinct_values`). This is compared with the number of stored nodes (`mean_traceback_size`). diff --git a/experiments/thresholding/thresholding.py b/experiments/thresholding/thresholding.py index c1b7aef..8f91c2d 100644 --- a/experiments/thresholding/thresholding.py +++ b/experiments/thresholding/thresholding.py @@ -206,7 +206,7 @@ def run_threading_task(num_samples, seed, quick): ts, h, true_parent, near = threading_case(num_samples, seed) m = ts.num_sites mu, rho = v.hmm_params(v.THREADING_DEFAULTS["k"]) - mi = tsls.MatcherIndexes(ts, alleles=v.ALLELES) + mi = tsls.TreeSequence(ts) true_switch = np.where(true_parent[1:] != true_parent[:-1])[0] + 1 rows = [] exact_posterior = None @@ -218,7 +218,7 @@ def run_threading_task(num_samples, seed, quick): rho, weight_by_n=False, count_distinct=True, - matcher_indexes=mi, + tree_sequence=mi, **kwargs, ) left, right, parent, path, matched, seconds = run_viterbi(am, ts, h) @@ -291,14 +291,14 @@ def run_calibration_task(setting, seed, quick): m = ts.num_sites n = ts.num_nodes rho = switches / m - mi = tsls.MatcherIndexes(ts, alleles=v.ALLELES) + mi = tsls.TreeSequence(ts) queries = [] for _ in range(3 if quick else 5): path = v.sample_prior_path(present, continuous, np.full(m, rho), n, rng) queries.append((path, v.mutate_along_path(A, path, mu, rng))) rows = [] for method, parameter, kwargs in method_grid(quick, rho, n, viterbi=False): - am = v.make_matcher(ts, mu, rho, weight_by_n=True, matcher_indexes=mi, **kwargs) + am = v.make_matcher(ts, mu, rho, weight_by_n=True, tree_sequence=mi, **kwargs) num_bins = len(CALIBRATION_BINS) - 1 count = np.zeros(num_bins) total_p = np.zeros(num_bins) @@ -353,7 +353,7 @@ def run_calibration_task(setting, seed, quick): @functools.lru_cache(maxsize=1) def load_sc2ts(size): """ - Return the (ts, matcher_indexes, H, kind, k) tuple for the sc2ts sub-ARG + Return the (ts, tree_sequence, H, kind, k) tuple for the sc2ts sub-ARG of the specified size, with the vestigial 0 -> 1 edge removed so that node 0 is isolated (as in benchmarking/engines.py). """ @@ -363,7 +363,7 @@ def load_sc2ts(size): tables = ts.dump_tables() tables.edges.keep_rows(tables.edges.parent != 0) ts = tables.tree_sequence() - mi = tsls.MatcherIndexes(ts, alleles=v.ALLELES) + mi = tsls.TreeSequence(ts) query_dir = SC2TS_DATA / "queries_scaling" H = np.load(query_dir / "queries.npz")["H"] df = pd.read_csv(query_dir / "queries.csv") @@ -421,7 +421,7 @@ def run_sc2ts_task(size, query_id, quick): rho, weight_by_n=False, count_distinct=True, - matcher_indexes=mi, + tree_sequence=mi, **kwargs, ) left, right, parent, _, matched, seconds = run_viterbi(am, ts, h) diff --git a/tsls/_tslsmodule.c b/tsls/_tslsmodule.c index c35b79c..37b6b49 100644 --- a/tsls/_tslsmodule.c +++ b/tsls/_tslsmodule.c @@ -23,14 +23,14 @@ static PyObject *TslsMatchImpossible; typedef struct { PyObject_HEAD - matcher_indexes_t *matcher_indexes; -} MatcherIndexes; + tsk_treeseq_t *tree_sequence; +} TreeSequence; typedef struct { PyObject_HEAD - ancestor_matcher_t *ancestor_matcher; - MatcherIndexes *matcher_indexes; -} AncestorMatcher; + tsls_hmm_t *hmm; + TreeSequence *tree_sequence; +} Hmm; static void handle_library_error(int err) @@ -46,224 +46,138 @@ handle_library_error(int err) } } -static FILE * -make_file(PyObject *fileobj, const char *mode) -{ - FILE *ret = NULL; - FILE *file = NULL; - int fileobj_fd, new_fd; +/*=================================================================== + * TreeSequence + *=================================================================== + */ - fileobj_fd = PyObject_AsFileDescriptor(fileobj); - if (fileobj_fd == -1) { - goto out; - } - new_fd = dup(fileobj_fd); - if (new_fd == -1) { - PyErr_SetFromErrno(PyExc_OSError); - goto out; - } - file = fdopen(new_fd, mode); - if (file == NULL) { - (void) close(new_fd); - PyErr_SetFromErrno(PyExc_OSError); - goto out; +/* A minimal tree sequence type, built from a LightweightTableCollection, + * which can be shared between HMMs (including in different threads). */ + +static int +TreeSequence_check_state(TreeSequence *self) +{ + int ret = 0; + if (self->tree_sequence == NULL) { + PyErr_SetString(PyExc_SystemError, "TreeSequence not initialised"); + ret = -1; } - ret = file; -out: return ret; } -/*=================================================================== - * MatcherIndexes - *=================================================================== - */ - static void -MatcherIndexes_dealloc(MatcherIndexes *self) +TreeSequence_dealloc(TreeSequence *self) { - if (self->matcher_indexes != NULL) { - matcher_indexes_free(self->matcher_indexes); - PyMem_Free(self->matcher_indexes); - self->matcher_indexes = NULL; + if (self->tree_sequence != NULL) { + tsk_treeseq_free(self->tree_sequence); + PyMem_Free(self->tree_sequence); + self->tree_sequence = NULL; } Py_TYPE(self)->tp_free((PyObject *) self); } static int -MatcherIndexes_init(MatcherIndexes *self, PyObject *args, PyObject *kwds) +TreeSequence_init(TreeSequence *self, PyObject *args, PyObject *kwds) { int ret = -1; int err; - LightweightTableCollection *tables; - PyObject *num_alleles_obj = NULL; - PyObject *ancestral_state_obj = NULL; - PyObject *derived_state_obj = NULL; - PyArrayObject *num_alleles_array = NULL; - PyArrayObject *ancestral_state_array = NULL; - PyArrayObject *derived_state_array = NULL; - tsk_size_t *num_alleles_data = NULL; - allele_t *ancestral_state_data = NULL; - allele_t *derived_state_data = NULL; - static char *kwlist[] - = { "tables", "num_alleles", "ancestral_state", "derived_state", NULL }; - - self->matcher_indexes = NULL; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!|OOO", kwlist, - &LightweightTableCollectionType, &tables, &num_alleles_obj, - &ancestral_state_obj, &derived_state_obj)) { + LightweightTableCollection *tables = NULL; + static char *kwlist[] = { "tables", NULL }; + + self->tree_sequence = NULL; + if (!PyArg_ParseTupleAndKeywords( + args, kwds, "O!", kwlist, &LightweightTableCollectionType, &tables)) { goto out; } if (LightweightTableCollection_check_state(tables) != 0) { goto out; } - if (num_alleles_obj != NULL && num_alleles_obj != Py_None) { - num_alleles_array = (PyArrayObject *) PyArray_FROMANY( - num_alleles_obj, NPY_UINT64, 1, 1, NPY_ARRAY_IN_ARRAY); - if (num_alleles_array == NULL) { - goto out; - } - if (PyArray_DIM(num_alleles_array, 0) - != (npy_intp) tables->tables->sites.num_rows) { - PyErr_SetString(PyExc_ValueError, "num_alleles must have length num_sites"); - goto out; - } - num_alleles_data = PyArray_DATA(num_alleles_array); - } - if (ancestral_state_obj != NULL && ancestral_state_obj != Py_None) { - ancestral_state_array = (PyArrayObject *) PyArray_FROMANY( - ancestral_state_obj, NPY_INT8, 1, 1, NPY_ARRAY_IN_ARRAY); - if (ancestral_state_array == NULL) { - goto out; - } - if (PyArray_DIM(ancestral_state_array, 0) - != (npy_intp) tables->tables->sites.num_rows) { - PyErr_SetString( - PyExc_ValueError, "ancestral_state must have length num_sites"); - goto out; - } - ancestral_state_data = PyArray_DATA(ancestral_state_array); - } - if (derived_state_obj != NULL && derived_state_obj != Py_None) { - derived_state_array = (PyArrayObject *) PyArray_FROMANY( - derived_state_obj, NPY_INT8, 1, 1, NPY_ARRAY_IN_ARRAY); - if (derived_state_array == NULL) { - goto out; - } - if (PyArray_DIM(derived_state_array, 0) - != (npy_intp) tables->tables->mutations.num_rows) { - PyErr_SetString( - PyExc_ValueError, "derived_state must have length num_mutations"); - goto out; - } - derived_state_data = PyArray_DATA(derived_state_array); - } - - self->matcher_indexes = PyMem_Calloc(1, sizeof(*self->matcher_indexes)); - if (self->matcher_indexes == NULL) { + self->tree_sequence = PyMem_Calloc(1, sizeof(*self->tree_sequence)); + if (self->tree_sequence == NULL) { PyErr_NoMemory(); goto out; } - err = matcher_indexes_alloc(self->matcher_indexes, tables->tables, - num_alleles_data, ancestral_state_data, derived_state_data, 0); + err = tsk_treeseq_init( + self->tree_sequence, tables->tables, TSK_TS_INIT_BUILD_INDEXES); if (err != 0) { - handle_library_error(err); + handle_tskit_error(err); goto out; } ret = 0; out: - Py_XDECREF(num_alleles_array); - Py_XDECREF(ancestral_state_array); - Py_XDECREF(derived_state_array); return ret; } -static int -MatcherIndexes_check_state(MatcherIndexes *self) +static PyObject * +TreeSequence_get_num_nodes(TreeSequence *self, void *closure) { - int ret = 0; - if (self->matcher_indexes == NULL) { - PyErr_SetString(PyExc_SystemError, "MatcherIndexes not initialised"); - ret = -1; + if (TreeSequence_check_state(self) != 0) { + return NULL; } - return ret; + return Py_BuildValue( + "n", (Py_ssize_t) tsk_treeseq_get_num_nodes(self->tree_sequence)); } static PyObject * -MatcherIndexes_print_state(MatcherIndexes *self, PyObject *args) +TreeSequence_get_num_sites(TreeSequence *self, void *closure) { - PyObject *ret = NULL; - PyObject *fileobj; - FILE *file = NULL; - - if (MatcherIndexes_check_state(self) != 0) { - goto out; + if (TreeSequence_check_state(self) != 0) { + return NULL; } - if (!PyArg_ParseTuple(args, "O", &fileobj)) { - goto out; - } - file = make_file(fileobj, "w"); - if (file == NULL) { - goto out; - } - matcher_indexes_print_state(self->matcher_indexes, file); - ret = Py_BuildValue(""); -out: - if (file != NULL) { - (void) fclose(file); - } - return ret; + return Py_BuildValue( + "n", (Py_ssize_t) tsk_treeseq_get_num_sites(self->tree_sequence)); } -static PyMethodDef MatcherIndexes_methods[] = { - { "print_state", (PyCFunction) MatcherIndexes_print_state, METH_VARARGS, - "Low-level debug method" }, +static PyGetSetDef TreeSequence_getsetters[] = { + { "num_nodes", (getter) TreeSequence_get_num_nodes, NULL, "The number of nodes" }, + { "num_sites", (getter) TreeSequence_get_num_sites, NULL, "The number of sites" }, { NULL } /* Sentinel */ }; -static PyTypeObject MatcherIndexesType = { +static PyTypeObject TreeSequenceType = { // clang-format off PyVarObject_HEAD_INIT(NULL, 0) - .tp_name = "_tsls.MatcherIndexes", - .tp_basicsize = sizeof(MatcherIndexes), - .tp_dealloc = (destructor) MatcherIndexes_dealloc, + .tp_name = "_tsls.TreeSequence", + .tp_basicsize = sizeof(TreeSequence), + .tp_dealloc = (destructor) TreeSequence_dealloc, .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, - .tp_doc = "MatcherIndexes objects", - .tp_methods = MatcherIndexes_methods, - .tp_init = (initproc) MatcherIndexes_init, + .tp_doc = "TreeSequence objects", + .tp_getset = TreeSequence_getsetters, + .tp_init = (initproc) TreeSequence_init, .tp_new = PyType_GenericNew, // clang-format on }; /*=================================================================== - * AncestorMatcher + * Hmm *=================================================================== */ static int -AncestorMatcher_check_state(AncestorMatcher *self) +Hmm_check_state(Hmm *self) { int ret = 0; - if (self->ancestor_matcher == NULL) { - PyErr_SetString(PyExc_SystemError, "AncestorMatcher not initialised"); + if (self->hmm == NULL) { + PyErr_SetString(PyExc_SystemError, "Hmm not initialised"); ret = -1; } return ret; } static void -AncestorMatcher_dealloc(AncestorMatcher *self) +Hmm_dealloc(Hmm *self) { - if (self->ancestor_matcher != NULL) { - ancestor_matcher_free(self->ancestor_matcher); - PyMem_Free(self->ancestor_matcher); - self->ancestor_matcher = NULL; + if (self->hmm != NULL) { + tsls_hmm_free(self->hmm); + PyMem_Free(self->hmm); + self->hmm = NULL; } - Py_XDECREF(self->matcher_indexes); + Py_XDECREF(self->tree_sequence); Py_TYPE(self)->tp_free((PyObject *) self); } static int -AncestorMatcher_init(AncestorMatcher *self, PyObject *args, PyObject *kwds) +Hmm_init(Hmm *self, PyObject *args, PyObject *kwds) { int ret = -1; int err; @@ -272,31 +186,59 @@ AncestorMatcher_init(AncestorMatcher *self, PyObject *args, PyObject *kwds) int pruning_mode = TSLS_PRUNE_NONE; double pruning_parameter = 0; int count_distinct = 0; - static char *kwlist[] = { "matcher_indexes", "recombination", "mismatch", - "extended_checks", "likelihood_threshold", "weight_by_n", "pruning_mode", - "pruning_parameter", "count_distinct", NULL }; - MatcherIndexes *matcher_indexes = NULL; + static char *kwlist[] = { "tree_sequence", "recombination", "mismatch", + "alleles", "extended_checks", "likelihood_threshold", "weight_by_n", + "pruning_mode", "pruning_parameter", "count_distinct", NULL }; + TreeSequence *tree_sequence = NULL; PyObject *recombination = NULL; PyObject *mismatch = NULL; + PyObject *alleles = Py_None; + PyObject *alleles_seq = NULL; + const char **alleles_data = NULL; + tsk_size_t num_alleles = 0; PyArrayObject *recombination_array = NULL; PyArrayObject *mismatch_array = NULL; npy_intp *shape; double likelihood_threshold = DBL_MIN; int flags = 0; + Py_ssize_t j; + size_t num_sites; - self->ancestor_matcher = NULL; - self->matcher_indexes = NULL; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!OO|idiidi", kwlist, - &MatcherIndexesType, &matcher_indexes, &recombination, &mismatch, + self->hmm = NULL; + self->tree_sequence = NULL; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!OO|Oidiidi", kwlist, + &TreeSequenceType, &tree_sequence, &recombination, &mismatch, &alleles, &extended_checks, &likelihood_threshold, &weight_by_n, &pruning_mode, &pruning_parameter, &count_distinct)) { goto out; } - self->matcher_indexes = matcher_indexes; - Py_INCREF(self->matcher_indexes); - if (MatcherIndexes_check_state(self->matcher_indexes) != 0) { + self->tree_sequence = tree_sequence; + Py_INCREF(self->tree_sequence); + if (TreeSequence_check_state(self->tree_sequence) != 0) { goto out; } + num_sites = tsk_treeseq_get_num_sites(tree_sequence->tree_sequence); + + if (alleles != Py_None) { + alleles_seq = PySequence_Fast(alleles, "alleles must be a sequence"); + if (alleles_seq == NULL) { + goto out; + } + num_alleles = (tsk_size_t) PySequence_Fast_GET_SIZE(alleles_seq); + alleles_data = PyMem_Calloc(num_alleles + 1, sizeof(*alleles_data)); + if (alleles_data == NULL) { + PyErr_NoMemory(); + goto out; + } + for (j = 0; j < (Py_ssize_t) num_alleles; j++) { + /* The UTF-8 buffer lives as long as the string */ + alleles_data[j] + = PyUnicode_AsUTF8(PySequence_Fast_GET_ITEM(alleles_seq, j)); + if (alleles_data[j] == NULL) { + goto out; + } + } + } recombination_array = (PyArrayObject *) PyArray_FromAny(recombination, PyArray_DescrFromType(NPY_FLOAT64), 1, 1, NPY_ARRAY_IN_ARRAY, NULL); @@ -304,7 +246,7 @@ AncestorMatcher_init(AncestorMatcher *self, PyObject *args, PyObject *kwds) goto out; } shape = PyArray_DIMS(recombination_array); - if (shape[0] != (npy_intp) matcher_indexes->matcher_indexes->num_sites) { + if (shape[0] != (npy_intp) num_sites) { PyErr_SetString( PyExc_ValueError, "Size of recombination array must be num_sites"); goto out; @@ -315,13 +257,13 @@ AncestorMatcher_init(AncestorMatcher *self, PyObject *args, PyObject *kwds) goto out; } shape = PyArray_DIMS(mismatch_array); - if (shape[0] != (npy_intp) matcher_indexes->matcher_indexes->num_sites) { + if (shape[0] != (npy_intp) num_sites) { PyErr_SetString(PyExc_ValueError, "Size of mismatch array must be num_sites"); goto out; } - self->ancestor_matcher = PyMem_Malloc(sizeof(ancestor_matcher_t)); - if (self->ancestor_matcher == NULL) { + self->hmm = PyMem_Calloc(1, sizeof(tsls_hmm_t)); + if (self->hmm == NULL) { PyErr_NoMemory(); goto out; } @@ -334,28 +276,29 @@ AncestorMatcher_init(AncestorMatcher *self, PyObject *args, PyObject *kwds) if (count_distinct) { flags |= TSLS_COUNT_DISTINCT; } - err = ancestor_matcher_alloc(self->ancestor_matcher, - self->matcher_indexes->matcher_indexes, PyArray_DATA(recombination_array), - PyArray_DATA(mismatch_array), likelihood_threshold, flags); + err = tsls_hmm_alloc(self->hmm, self->tree_sequence->tree_sequence, alleles_data, + num_alleles, PyArray_DATA(recombination_array), PyArray_DATA(mismatch_array), + likelihood_threshold, flags); if (err != 0) { handle_library_error(err); goto out; } - err = ancestor_matcher_set_pruning( - self->ancestor_matcher, pruning_mode, pruning_parameter); + err = tsls_hmm_set_pruning(self->hmm, pruning_mode, pruning_parameter); if (err != 0) { handle_library_error(err); goto out; } ret = 0; out: + PyMem_Free(alleles_data); + Py_XDECREF(alleles_seq); Py_XDECREF(recombination_array); Py_XDECREF(mismatch_array); return ret; } static PyObject * -AncestorMatcher_find_path(AncestorMatcher *self, PyObject *args, PyObject *kwds) +Hmm_find_path(Hmm *self, PyObject *args, PyObject *kwds) { int err; PyObject *ret = NULL; @@ -371,7 +314,7 @@ AncestorMatcher_find_path(AncestorMatcher *self, PyObject *args, PyObject *kwds) PyArrayObject *match = NULL; npy_intp dims[1]; - if (AncestorMatcher_check_state(self) != 0) { + if (Hmm_check_state(self) != 0) { goto out; } if (!PyArg_ParseTupleAndKeywords( @@ -388,14 +331,14 @@ AncestorMatcher_find_path(AncestorMatcher *self, PyObject *args, PyObject *kwds) goto out; } shape = PyArray_DIMS(haplotype_array); - if (shape[0] != (npy_intp) self->ancestor_matcher->num_sites) { + if (shape[0] != (npy_intp) self->hmm->num_sites) { PyErr_SetString(PyExc_ValueError, "Incorrect size for input haplotype."); goto out; } - dims[0] = self->ancestor_matcher->num_sites; - left = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_UINT32); - right = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_UINT32); + dims[0] = self->hmm->num_sites; + left = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_FLOAT64); + right = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_FLOAT64); parent = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_INT32); match = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_INT8); if (left == NULL || right == NULL || parent == NULL || match == NULL) { @@ -403,7 +346,7 @@ AncestorMatcher_find_path(AncestorMatcher *self, PyObject *args, PyObject *kwds) } Py_BEGIN_ALLOW_THREADS - err = ancestor_matcher_find_path(self->ancestor_matcher, (tsk_id_t) start, + err = tsls_hmm_find_path(self->hmm, (tsk_id_t) start, (tsk_id_t) end, (allele_t *) PyArray_DATA(haplotype_array), (allele_t *) PyArray_DATA(match), &num_edges, PyArray_DATA(left), PyArray_DATA(right), PyArray_DATA(parent)); @@ -506,14 +449,14 @@ compressed_matrix_to_tuple(const compressed_matrix_t *cm, bool with_flags) * (normalisation_factor, offsets, nodes, values), where the values at site * j are nodes[offsets[j]: offsets[j + 1]] and similarly for values. */ static PyObject * -AncestorMatcher_run_compressed_matrix( - AncestorMatcher *self, PyObject *haplotype, PyObject *forward_norm) +Hmm_run_compressed_matrix( + Hmm *self, PyObject *haplotype, PyObject *forward_norm) { int err; PyObject *ret = NULL; PyArrayObject *haplotype_array = NULL; PyArrayObject *forward_norm_array = NULL; - const size_t num_sites = self->ancestor_matcher->num_sites; + const size_t num_sites = self->hmm->num_sites; compressed_matrix_t cm; bool cm_allocated = false; @@ -539,7 +482,7 @@ AncestorMatcher_run_compressed_matrix( } } err = compressed_matrix_alloc(&cm, num_sites, - 2 * (self->ancestor_matcher->num_nodes + 1) * sizeof(double)); + 2 * (self->hmm->num_nodes + 1) * sizeof(double)); cm_allocated = true; if (err != 0) { handle_library_error(err); @@ -548,10 +491,10 @@ AncestorMatcher_run_compressed_matrix( Py_BEGIN_ALLOW_THREADS if (forward_norm_array == NULL) { - err = ancestor_matcher_forward(self->ancestor_matcher, + err = tsls_hmm_forward(self->hmm, (allele_t *) PyArray_DATA(haplotype_array), &cm); } else { - err = ancestor_matcher_backward(self->ancestor_matcher, + err = tsls_hmm_backward(self->hmm, (allele_t *) PyArray_DATA(haplotype_array), (double *) PyArray_DATA(forward_norm_array), &cm); } @@ -574,7 +517,7 @@ AncestorMatcher_run_compressed_matrix( /* Run the Viterbi algorithm, returning the find_path outputs followed by * the compressed matrix of normalised Viterbi values with flags. */ static PyObject * -AncestorMatcher_viterbi_matrix(AncestorMatcher *self, PyObject *args, PyObject *kwds) +Hmm_viterbi_matrix(Hmm *self, PyObject *args, PyObject *kwds) { int err; PyObject *ret = NULL; @@ -593,14 +536,14 @@ AncestorMatcher_viterbi_matrix(AncestorMatcher *self, PyObject *args, PyObject * npy_intp dims[1]; size_t num_sites; - if (AncestorMatcher_check_state(self) != 0) { + if (Hmm_check_state(self) != 0) { goto out; } if (!PyArg_ParseTupleAndKeywords( args, kwds, "Oii", kwlist, &haplotype, &start, &end)) { goto out; } - num_sites = self->ancestor_matcher->num_sites; + num_sites = self->hmm->num_sites; haplotype_array = (PyArrayObject *) PyArray_FROMANY(haplotype, NPY_INT8, 1, 1, NPY_ARRAY_IN_ARRAY); if (haplotype_array == NULL) { @@ -611,22 +554,22 @@ AncestorMatcher_viterbi_matrix(AncestorMatcher *self, PyObject *args, PyObject * goto out; } dims[0] = (npy_intp) num_sites; - left = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_UINT32); - right = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_UINT32); + left = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_FLOAT64); + right = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_FLOAT64); parent = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_INT32); match = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_INT8); if (left == NULL || right == NULL || parent == NULL || match == NULL) { goto out; } err = compressed_matrix_alloc( - &cm, num_sites, 2 * (self->ancestor_matcher->num_nodes + 1) * sizeof(double)); + &cm, num_sites, 2 * (self->hmm->num_nodes + 1) * sizeof(double)); cm_allocated = true; if (err != 0) { handle_library_error(err); goto out; } Py_BEGIN_ALLOW_THREADS - err = ancestor_matcher_find_path_matrix(self->ancestor_matcher, (tsk_id_t) start, + err = tsls_hmm_find_path_matrix(self->hmm, (tsk_id_t) start, (tsk_id_t) end, (allele_t *) PyArray_DATA(haplotype_array), (allele_t *) PyArray_DATA(match), &num_edges, PyArray_DATA(left), PyArray_DATA(right), PyArray_DATA(parent), &cm); @@ -663,45 +606,45 @@ AncestorMatcher_viterbi_matrix(AncestorMatcher *self, PyObject *args, PyObject * } static PyObject * -AncestorMatcher_forward_matrix(AncestorMatcher *self, PyObject *args, PyObject *kwds) +Hmm_forward_matrix(Hmm *self, PyObject *args, PyObject *kwds) { PyObject *ret = NULL; static char *kwlist[] = { "haplotype", NULL }; PyObject *haplotype = NULL; - if (AncestorMatcher_check_state(self) != 0) { + if (Hmm_check_state(self) != 0) { goto out; } if (!PyArg_ParseTupleAndKeywords(args, kwds, "O", kwlist, &haplotype)) { goto out; } - ret = AncestorMatcher_run_compressed_matrix(self, haplotype, NULL); + ret = Hmm_run_compressed_matrix(self, haplotype, NULL); out: return ret; } static PyObject * -AncestorMatcher_backward_matrix(AncestorMatcher *self, PyObject *args, PyObject *kwds) +Hmm_backward_matrix(Hmm *self, PyObject *args, PyObject *kwds) { PyObject *ret = NULL; static char *kwlist[] = { "haplotype", "normalisation_factor", NULL }; PyObject *haplotype = NULL; PyObject *forward_norm = NULL; - if (AncestorMatcher_check_state(self) != 0) { + if (Hmm_check_state(self) != 0) { goto out; } if (!PyArg_ParseTupleAndKeywords( args, kwds, "OO", kwlist, &haplotype, &forward_norm)) { goto out; } - ret = AncestorMatcher_run_compressed_matrix(self, haplotype, forward_norm); + ret = Hmm_run_compressed_matrix(self, haplotype, forward_norm); out: return ret; } static PyObject * -AncestorMatcher_get_traceback(AncestorMatcher *self, PyObject *args) +Hmm_get_traceback(Hmm *self, PyObject *args) { PyObject *ret = NULL; unsigned long site; @@ -711,13 +654,13 @@ AncestorMatcher_get_traceback(AncestorMatcher *self, PyObject *args) PyObject *value = NULL; int j; - if (AncestorMatcher_check_state(self) != 0) { + if (Hmm_check_state(self) != 0) { goto out; } if (!PyArg_ParseTuple(args, "k", &site)) { goto out; } - if (site >= self->ancestor_matcher->num_sites) { + if (site >= self->hmm->num_sites) { PyErr_SetString(PyExc_ValueError, "site out of range"); goto out; } @@ -725,7 +668,7 @@ AncestorMatcher_get_traceback(AncestorMatcher *self, PyObject *args) if (dict == NULL) { goto out; } - list = &self->ancestor_matcher->traceback[site]; + list = &self->hmm->traceback[site]; for (j = 0; j < list->size; j++) { key = Py_BuildValue("k", (unsigned long) list->node[j]); value = Py_BuildValue("i", (int) list->recombination_required[j]); @@ -750,40 +693,40 @@ AncestorMatcher_get_traceback(AncestorMatcher *self, PyObject *args) } static PyObject * -AncestorMatcher_get_mean_traceback_size(AncestorMatcher *self, void *closure) +Hmm_get_mean_traceback_size(Hmm *self, void *closure) { PyObject *ret = NULL; - if (AncestorMatcher_check_state(self) != 0) { + if (Hmm_check_state(self) != 0) { goto out; } ret = Py_BuildValue( - "d", ancestor_matcher_get_mean_traceback_size(self->ancestor_matcher)); + "d", tsls_hmm_get_mean_traceback_size(self->hmm)); out: return ret; } static PyObject * -AncestorMatcher_get_mean_num_distinct_values(AncestorMatcher *self, void *closure) +Hmm_get_mean_num_distinct_values(Hmm *self, void *closure) { PyObject *ret = NULL; - if (AncestorMatcher_check_state(self) != 0) { + if (Hmm_check_state(self) != 0) { goto out; } ret = Py_BuildValue( - "d", ancestor_matcher_get_mean_num_distinct_values(self->ancestor_matcher)); + "d", tsls_hmm_get_mean_num_distinct_values(self->hmm)); out: return ret; } static PyObject * -AncestorMatcher_get_total_memory(AncestorMatcher *self, void *closure) +Hmm_get_total_memory(Hmm *self, void *closure) { PyObject *ret = NULL; unsigned long val; - if (AncestorMatcher_check_state(self) != 0) { + if (Hmm_check_state(self) != 0) { goto out; } @@ -791,7 +734,7 @@ AncestorMatcher_get_total_memory(AncestorMatcher *self, void *closure) /* Without atomics, return an obviously wrong value */ val = (unsigned long) PY_SSIZE_T_MAX; #else - val = ancestor_matcher_get_total_memory(self->ancestor_matcher); + val = tsls_hmm_get_total_memory(self->hmm); #endif ret = Py_BuildValue("k", val); @@ -799,47 +742,47 @@ AncestorMatcher_get_total_memory(AncestorMatcher *self, void *closure) return ret; } -static PyGetSetDef AncestorMatcher_getsetters[] = { - { "mean_traceback_size", (getter) AncestorMatcher_get_mean_traceback_size, NULL, +static PyGetSetDef Hmm_getsetters[] = { + { "mean_traceback_size", (getter) Hmm_get_mean_traceback_size, NULL, "The mean size of the traceback per site." }, { "mean_num_distinct_values", - (getter) AncestorMatcher_get_mean_num_distinct_values, NULL, + (getter) Hmm_get_mean_num_distinct_values, NULL, "The mean number of distinct likelihood values per site in the last " "Viterbi run (if count_distinct is set)." }, - { "total_memory", (getter) AncestorMatcher_get_total_memory, NULL, - "The total amount of memory used by this matcher." }, + { "total_memory", (getter) Hmm_get_total_memory, NULL, + "The total amount of memory used by this HMM." }, { NULL } /* Sentinel */ }; -static PyMethodDef AncestorMatcher_methods[] = { - { "find_path", (PyCFunction) AncestorMatcher_find_path, METH_VARARGS | METH_KEYWORDS, +static PyMethodDef Hmm_methods[] = { + { "find_path", (PyCFunction) Hmm_find_path, METH_VARARGS | METH_KEYWORDS, "Returns a best match path for the specified haplotype through the ancestors." }, - { "viterbi_matrix", (PyCFunction) AncestorMatcher_viterbi_matrix, + { "viterbi_matrix", (PyCFunction) Hmm_viterbi_matrix, METH_VARARGS | METH_KEYWORDS, "Runs the Viterbi algorithm and also returns the compressed matrix of " "normalised values and recombination flags." }, - { "forward_matrix", (PyCFunction) AncestorMatcher_forward_matrix, + { "forward_matrix", (PyCFunction) Hmm_forward_matrix, METH_VARARGS | METH_KEYWORDS, "Runs the forward algorithm and returns the compressed matrix arrays." }, - { "backward_matrix", (PyCFunction) AncestorMatcher_backward_matrix, + { "backward_matrix", (PyCFunction) Hmm_backward_matrix, METH_VARARGS | METH_KEYWORDS, "Runs the backward algorithm and returns the compressed matrix arrays." }, - { "get_traceback", (PyCFunction) AncestorMatcher_get_traceback, METH_VARARGS, + { "get_traceback", (PyCFunction) Hmm_get_traceback, METH_VARARGS, "Returns the traceback likelihood dictionary at the specified site." }, { NULL } /* Sentinel */ }; -static PyTypeObject AncestorMatcherType = { +static PyTypeObject HmmType = { // clang-format off PyVarObject_HEAD_INIT(NULL, 0) - .tp_name = "_tsls.AncestorMatcher", - .tp_basicsize = sizeof(AncestorMatcher), - .tp_dealloc = (destructor) AncestorMatcher_dealloc, + .tp_name = "_tsls.Hmm", + .tp_basicsize = sizeof(Hmm), + .tp_dealloc = (destructor) Hmm_dealloc, .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, - .tp_doc = "AncestorMatcher objects", - .tp_methods = AncestorMatcher_methods, - .tp_getset = AncestorMatcher_getsetters, - .tp_init = (initproc) AncestorMatcher_init, + .tp_doc = "Hmm objects", + .tp_methods = Hmm_methods, + .tp_getset = Hmm_getsetters, + .tp_init = (initproc) Hmm_init, .tp_new = PyType_GenericNew, // clang-format on }; @@ -891,21 +834,21 @@ init_tsls(void) register_lwt_class(module); - /* MatcherIndexes type */ - MatcherIndexesType.tp_new = PyType_GenericNew; - if (PyType_Ready(&MatcherIndexesType) < 0) { + /* TreeSequence type */ + TreeSequenceType.tp_new = PyType_GenericNew; + if (PyType_Ready(&TreeSequenceType) < 0) { INITERROR; } - Py_INCREF(&MatcherIndexesType); - PyModule_AddObject(module, "MatcherIndexes", (PyObject *) &MatcherIndexesType); + Py_INCREF(&TreeSequenceType); + PyModule_AddObject(module, "TreeSequence", (PyObject *) &TreeSequenceType); - /* AncestorMatcher type */ - AncestorMatcherType.tp_new = PyType_GenericNew; - if (PyType_Ready(&AncestorMatcherType) < 0) { + /* Hmm type */ + HmmType.tp_new = PyType_GenericNew; + if (PyType_Ready(&HmmType) < 0) { INITERROR; } - Py_INCREF(&AncestorMatcherType); - PyModule_AddObject(module, "AncestorMatcher", (PyObject *) &AncestorMatcherType); + Py_INCREF(&HmmType); + PyModule_AddObject(module, "Hmm", (PyObject *) &HmmType); TslsLibraryError = PyErr_NewException("_tsls.LibraryError", NULL, NULL); Py_INCREF(TslsLibraryError); diff --git a/tsls/lib/err.c b/tsls/lib/err.c index ba426d9..3581cb5 100644 --- a/tsls/lib/err.c +++ b/tsls/lib/err.c @@ -30,23 +30,11 @@ tsls_strerror(int err) case TSLS_ERR_BAD_HAPLOTYPE_ALLELE: ret = "Haplotype allele must be missing (-1) or >= 0 and < num_alleles"; break; - case TSLS_ERR_BAD_ANCESTRAL_STATE: - ret = "Ancestral state must be >= 0 and < num_alleles"; + case TSLS_ERR_UNKNOWN_ALLELE: + ret = "An ancestral or derived state is not in the list of alleles"; break; - case TSLS_ERR_BAD_MUTATION_DERIVED_STATE: - ret = "Mutation derived state must be >= 0 and < num_alleles"; - break; - case TSLS_ERR_BAD_MUTATION_NODE: - ret = "Bad mutation: node out of bounds"; - break; - case TSLS_ERR_BAD_MUTATION_SITE: - ret = "Bad mutation: site out of bounds"; - break; - case TSLS_ERR_BAD_EDGE_NODE: - ret = "Bad edge: parent or child node out of bounds"; - break; - case TSLS_ERR_BAD_COORDINATE: - ret = "Edge and site coordinates must be integers"; + case TSLS_ERR_TOO_MANY_ALLELES: + ret = "Too many distinct alleles at a site"; break; case TSLS_ERR_BAD_SITE_RANGE: ret = "Bad site range: must have 0 <= start < end <= num_sites"; diff --git a/tsls/lib/err.h b/tsls/lib/err.h index d37c03b..e1c811e 100644 --- a/tsls/lib/err.h +++ b/tsls/lib/err.h @@ -8,15 +8,11 @@ #define TSLS_ERR_MATCH_IMPOSSIBLE_EXTREME_MUTATION_PROBA -4 #define TSLS_ERR_MATCH_IMPOSSIBLE_ZERO_RECOMB_PRECISION -5 #define TSLS_ERR_BAD_HAPLOTYPE_ALLELE -6 -#define TSLS_ERR_BAD_ANCESTRAL_STATE -7 -#define TSLS_ERR_BAD_MUTATION_DERIVED_STATE -8 -#define TSLS_ERR_BAD_MUTATION_NODE -9 -#define TSLS_ERR_BAD_MUTATION_SITE -10 -#define TSLS_ERR_BAD_EDGE_NODE -11 -#define TSLS_ERR_BAD_COORDINATE -12 -#define TSLS_ERR_BAD_SITE_RANGE -13 -#define TSLS_ERR_BAD_NORMALISATION_FACTOR -14 -#define TSLS_ERR_BAD_PRUNING -15 +#define TSLS_ERR_UNKNOWN_ALLELE -7 +#define TSLS_ERR_TOO_MANY_ALLELES -8 +#define TSLS_ERR_BAD_SITE_RANGE -9 +#define TSLS_ERR_BAD_NORMALISATION_FACTOR -10 +#define TSLS_ERR_BAD_PRUNING -11 // clang-format on #ifdef __GNUC__ diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index 92f1105..73f8336 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -38,12 +38,10 @@ #define MODE_BACKWARD 2 static inline void -remove_edge(edge_t edge, tsk_id_t *restrict parent, tsk_id_t *restrict left_child, - tsk_id_t *restrict right_child, tsk_id_t *restrict left_sib, - tsk_id_t *restrict right_sib) +remove_edge(tsk_id_t p, tsk_id_t c, tsk_id_t *restrict parent, + tsk_id_t *restrict left_child, tsk_id_t *restrict right_child, + tsk_id_t *restrict left_sib, tsk_id_t *restrict right_sib) { - const tsk_id_t p = edge.parent; - const tsk_id_t c = edge.child; const tsk_id_t lsib = left_sib[c]; const tsk_id_t rsib = right_sib[c]; @@ -63,12 +61,10 @@ remove_edge(edge_t edge, tsk_id_t *restrict parent, tsk_id_t *restrict left_chil } static inline void -insert_edge(edge_t edge, tsk_id_t *restrict parent, tsk_id_t *restrict left_child, - tsk_id_t *restrict right_child, tsk_id_t *restrict left_sib, - tsk_id_t *restrict right_sib) +insert_edge(tsk_id_t p, tsk_id_t c, tsk_id_t *restrict parent, + tsk_id_t *restrict left_child, tsk_id_t *restrict right_child, + tsk_id_t *restrict left_sib, tsk_id_t *restrict right_sib) { - const tsk_id_t p = edge.parent; - const tsk_id_t c = edge.child; const tsk_id_t u = right_child[p]; parent[c] = p; @@ -99,265 +95,33 @@ 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) -{ - return (double) (coordinate_t) x == x; -} - -/*=================================================================== - * Matcher indexes - *=================================================================== - */ - -int -matcher_indexes_print_state(const matcher_indexes_t *self, FILE *out) -{ - size_t j; - mutation_list_node_t *u; - - fprintf(out, "Matcher indexes state\n"); - fprintf(out, "flags = %d\n", (int) self->flags); - fprintf(out, "num_sites = %d\n", (int) self->num_sites); - fprintf(out, "num_nodes = %d\n", (int) self->num_nodes); - fprintf(out, "num_edges = %d\n", (int) self->num_edges); - - fprintf(out, "Sites = \n"); - fprintf(out, "site\tposition\tancestral_state\tnum_alleles\t(node, derived_state)\n"); - for (j = 0; j < self->num_sites; j++) { - fprintf(out, "%d\t%d\t%d\t%d\t", (int) j, (int) self->sites.position[j], - (int) self->sites.ancestral_state[j], (int) self->sites.num_alleles[j]); - for (u = self->sites.mutations[j]; u != NULL; u = u->next) { - fprintf(out, "(%d, %d) ", u->node, u->derived_state); - } - fprintf(out, "\n"); - } - return 0; -} - -static void -matcher_indexes_copy_edge_indexes( - matcher_indexes_t *self, const tsk_table_collection_t *tables) -{ - tsk_size_t j; - tsk_id_t k; - edge_t e; - const tsk_id_t *restrict I = tables->indexes.edge_insertion_order; - const tsk_id_t *restrict O = tables->indexes.edge_removal_order; - const double *restrict edges_right = tables->edges.right; - const double *restrict edges_left = tables->edges.left; - const tsk_id_t *restrict edges_child = tables->edges.child; - const tsk_id_t *restrict edges_parent = tables->edges.parent; - - for (j = 0; j < self->num_edges; j++) { - k = I[j]; - e.left = (coordinate_t) edges_left[k]; - e.right = (coordinate_t) edges_right[k]; - e.parent = edges_parent[k]; - e.child = edges_child[k]; - self->left_index_edges[j] = e; - - k = O[j]; - e.left = (coordinate_t) edges_left[k]; - e.right = (coordinate_t) edges_right[k]; - e.parent = edges_parent[k]; - e.child = edges_child[k]; - self->right_index_edges[j] = e; - } -} - +/* Position the specified null tree position at the tree containing x, + * seeking from whichever end of the sequence is closer, as + * tsk_tree_seek_from_null does. The edges of the tree are then the edges + * in the tp->in range (stepping by tp->direction) that cover x. */ static int WARN_UNUSED -matcher_indexes_add_mutation( - matcher_indexes_t *self, tsk_id_t site, tsk_id_t node, allele_t derived_state) +seek_tree_position(tsk_tree_position_t *tp, double x) { int ret = 0; - mutation_list_node_t *list_node, *tail; - - list_node = tsi_blkalloc_get(&self->allocator, sizeof(mutation_list_node_t)); - if (list_node == NULL) { - ret = TSLS_ERR_NO_MEMORY; - goto out; - } - list_node->node = node; - list_node->derived_state = derived_state; - list_node->next = NULL; - if (self->sites.mutations[site] == NULL) { - self->sites.mutations[site] = list_node; + const tsk_treeseq_t *ts = tp->tree_sequence; + const double *restrict breakpoints = ts->breakpoints; + const tsk_size_t num_trees = ts->num_trees; + tsk_id_t index; + + assert(tp->index == -1); + assert(0 <= x && x < ts->tables->sequence_length); + index = (tsk_id_t) tsk_search_sorted(breakpoints, num_trees + 1, x); + if (breakpoints[index] > x) { + index--; + } + if (x <= ts->tables->sequence_length / 2) { + ret = tsk_tree_position_seek_forward(tp, index); } else { - tail = self->sites.mutations[site]; - while (tail->next != NULL) { - tail = tail->next; - } - tail->next = list_node; + ret = tsk_tree_position_seek_backward(tp, index); } - self->num_mutations++; -out: return ret; } -/* Copy the site and mutation data into the matcher indexes. The - * ancestral_state (per site) and derived_state (per mutation) arrays - * give the integer-encoded alleles; if NULL, the ancestral state is 0 - * and every mutation has derived state 1. Mutations must be in tskit - * order, so that when two mutations at a site are on the same node the - * later one (the child) takes precedence. */ -static int -matcher_indexes_copy_mutation_data(matcher_indexes_t *self, - const tsk_table_collection_t *tables, const tsk_size_t *num_alleles, - const allele_t *ancestral_state, const allele_t *derived_state) -{ - int ret = 0; - tsk_size_t j; - tsk_id_t site; - allele_t state; - const double *restrict sites_position = tables->sites.position; - const tsk_id_t *restrict mutations_site = tables->mutations.site; - const tsk_id_t *restrict mutations_node = tables->mutations.node; - const tsk_size_t total_mutations = tables->mutations.num_rows; - coordinate_t *restrict converted_position = self->sites.position; - - for (j = 0; j < self->num_sites; j++) { - converted_position[j] = (coordinate_t) sites_position[j]; - self->sites.mutations[j] = NULL; - self->sites.num_alleles[j] = num_alleles != NULL ? num_alleles[j] : 2; - state = ancestral_state != NULL ? ancestral_state[j] : 0; - if (state < 0 || (tsk_size_t) state >= self->sites.num_alleles[j]) { - ret = TSLS_ERR_BAD_ANCESTRAL_STATE; - goto out; - } - self->sites.ancestral_state[j] = state; - } - converted_position[j] = (coordinate_t) tables->sequence_length; - - for (j = 0; j < total_mutations; j++) { - site = mutations_site[j]; - state = derived_state != NULL ? derived_state[j] : 1; - if (state < 0 || (tsk_size_t) state >= self->sites.num_alleles[site]) { - ret = TSLS_ERR_BAD_MUTATION_DERIVED_STATE; - goto out; - } - ret = matcher_indexes_add_mutation(self, site, mutations_node[j], state); - if (ret != 0) { - goto out; - } - } -out: - return ret; -} - -static int -matcher_indexes_validate_tables(const tsk_table_collection_t *tables) -{ - tsk_size_t j; - const tsk_id_t num_nodes = (tsk_id_t) tables->nodes.num_rows; - const tsk_id_t num_sites = (tsk_id_t) tables->sites.num_rows; - - if (!is_integer(tables->sequence_length)) { - return TSLS_ERR_BAD_COORDINATE; - } - for (j = 0; j < tables->edges.num_rows; j++) { - if (tables->edges.parent[j] < 0 || tables->edges.parent[j] >= num_nodes - || tables->edges.child[j] < 0 || tables->edges.child[j] >= num_nodes) { - return TSLS_ERR_BAD_EDGE_NODE; - } - if (!is_integer(tables->edges.left[j]) || !is_integer(tables->edges.right[j])) { - return TSLS_ERR_BAD_COORDINATE; - } - } - for (j = 0; j < tables->sites.num_rows; j++) { - if (!is_integer(tables->sites.position[j])) { - return TSLS_ERR_BAD_COORDINATE; - } - } - for (j = 0; j < tables->mutations.num_rows; j++) { - if (tables->mutations.node[j] < 0 || tables->mutations.node[j] >= num_nodes) { - return TSLS_ERR_BAD_MUTATION_NODE; - } - if (tables->mutations.site[j] < 0 || tables->mutations.site[j] >= num_sites) { - return TSLS_ERR_BAD_MUTATION_SITE; - } - } - return 0; -} - -/* The tables must be sorted and indexed, as in a valid tree sequence. */ -int -matcher_indexes_alloc(matcher_indexes_t *self, const tsk_table_collection_t *tables, - const tsk_size_t *num_alleles, const allele_t *ancestral_state, - const allele_t *derived_state, tsk_flags_t flags) -{ - int ret = 0; - - memset(self, 0, sizeof(*self)); - self->flags = flags; - self->num_edges = tables->edges.num_rows; - self->num_nodes = tables->nodes.num_rows; - self->num_sites = tables->sites.num_rows; - /* num_mutations is incremented as mutations are added */ - self->num_mutations = 0; - - ret = matcher_indexes_validate_tables(tables); - if (ret != 0) { - goto out; - } - - self->left_index_edges = malloc(self->num_edges * sizeof(*self->left_index_edges)); - self->right_index_edges = malloc(self->num_edges * sizeof(*self->right_index_edges)); - self->sites.mutations = malloc(self->num_sites * sizeof(*self->sites.mutations)); - self->sites.num_alleles = malloc(self->num_sites * sizeof(*self->sites.num_alleles)); - self->sites.ancestral_state - = malloc(self->num_sites * sizeof(*self->sites.ancestral_state)); - self->sites.position - = malloc((self->num_sites + 1) * sizeof(*self->sites.position)); - if (self->left_index_edges == NULL || self->right_index_edges == NULL - || self->sites.mutations == NULL || self->sites.position == NULL - || self->sites.num_alleles == NULL || self->sites.ancestral_state == NULL) { - ret = TSLS_ERR_NO_MEMORY; - goto out; - } - if (tsi_blkalloc_init(&self->allocator, 65536) != 0) { - ret = TSLS_ERR_NO_MEMORY; - goto out; - } - matcher_indexes_copy_edge_indexes(self, tables); - ret = matcher_indexes_copy_mutation_data( - self, tables, num_alleles, ancestral_state, derived_state); - if (ret != 0) { - goto out; - } -out: - return ret; -} - -int -matcher_indexes_free(matcher_indexes_t *self) -{ - tsi_safe_free(self->left_index_edges); - tsi_safe_free(self->right_index_edges); - tsi_safe_free(self->sites.mutations); - tsi_safe_free(self->sites.position); - tsi_safe_free(self->sites.num_alleles); - tsi_safe_free(self->sites.ancestral_state); - tsi_blkalloc_free(&self->allocator); - return 0; -} - /*=================================================================== * Compressed matrix *=================================================================== @@ -466,18 +230,18 @@ compressed_matrix_store_site(compressed_matrix_t *self, tsk_id_t site, } /*=================================================================== - * Ancestor matcher + * HMM *=================================================================== */ static inline bool -ancestor_matcher_in_tree(const ancestor_matcher_t *self, tsk_id_t u) +tsls_hmm_in_tree(const tsls_hmm_t *self, tsk_id_t u) { return self->parent[u] != NULL_NODE || self->left_child[u] != NULL_NODE; } static void -ancestor_matcher_check_state(const ancestor_matcher_t *self) +tsls_hmm_check_state(const tsls_hmm_t *self) { int num_likelihoods = 0; int j; @@ -491,7 +255,7 @@ ancestor_matcher_check_state(const ancestor_matcher_t *self) if (self->likelihood[u] >= 0) { num_likelihoods++; } - if (ancestor_matcher_in_tree(self, u)) { + if (tsls_hmm_in_tree(self, u)) { assert(self->likelihood[u] != NOT_IN_TREE); if (self->parent[u] == NULL_NODE) { /* Roots always have a likelihood */ @@ -508,16 +272,17 @@ ancestor_matcher_check_state(const ancestor_matcher_t *self) } int -ancestor_matcher_print_state(ancestor_matcher_t *self, FILE *out) +tsls_hmm_print_state(tsls_hmm_t *self, FILE *out) { int j, k; tsk_id_t u; - fprintf(out, "Ancestor matcher state\n"); - fprintf(out, "site\trecomb_rate\tmut_rate\n"); + fprintf(out, "HMM state\n"); + fprintf(out, "site\tposition\trecomb_rate\tmut_rate\tnum_alleles\tancestral\n"); for (j = 0; j < (int) self->num_sites; j++) { - fprintf( - out, "%d\t%f\t%f\n", j, self->recombination_rate[j], self->mismatch_rate[j]); + fprintf(out, "%d\t%f\t%f\t%f\t%d\t%d\n", j, self->sites_position[j], + self->recombination_rate[j], self->mismatch_rate[j], + (int) self->num_alleles[j], (int) self->ancestral_state[j]); } fprintf(out, "tree = \n"); fprintf(out, "id\tparent\tlchild\trchild\tlsib\trsib\tlikelihood\n"); @@ -545,20 +310,126 @@ ancestor_matcher_print_state(ancestor_matcher_t *self, FILE *out) return 0; } +/* Returns the index of the specified state in the list of alleles, or -1 + * if it is not in the list. */ +static int +find_allele(const char *state, tsk_size_t length, const char **alleles, + const tsk_size_t *allele_length, tsk_size_t num_alleles) +{ + tsk_size_t j; + + for (j = 0; j < num_alleles; j++) { + if (allele_length[j] == length && memcmp(alleles[j], state, length) == 0) { + return (int) j; + } + } + return -1; +} + +/* Returns the code for the specified state. If the state is not in the list + * of alleles and add is true, it is appended to the list. */ +static int +encode_allele(const char *state, tsk_size_t length, const char **alleles, + tsk_size_t *allele_length, tsk_size_t *num_alleles, bool add) +{ + int ret = find_allele(state, length, alleles, allele_length, *num_alleles); + + if (ret < 0) { + if (!add) { + ret = TSLS_ERR_UNKNOWN_ALLELE; + } else if (*num_alleles == TSLS_MAX_ALLELES) { + ret = TSLS_ERR_TOO_MANY_ALLELES; + } else { + alleles[*num_alleles] = state; + allele_length[*num_alleles] = length; + ret = (int) *num_alleles; + (*num_alleles)++; + } + } + return ret; +} + +/* Encode the ancestral and derived states of the tree sequence as integers. + * If alleles is NULL, the alleles are numbered at each site as in tskit's + * genotype decoding: the ancestral state is 0, and the derived states are + * numbered in order of their first appearance at the site, with at least + * two alleles at each site. Otherwise, each + * state is its index in the specified list of alleles. */ +static int WARN_UNUSED +tsls_hmm_encode_alleles( + tsls_hmm_t *self, const char **alleles, tsk_size_t num_global_alleles) +{ + int ret = 0; + const tsk_treeseq_t *ts = self->tree_sequence; + const tsk_site_table_t *sites = &ts->tables->sites; + const bool global = alleles != NULL; + const char *site_alleles[TSLS_MAX_ALLELES]; + tsk_size_t allele_length[TSLS_MAX_ALLELES]; + const char **list = global ? alleles : site_alleles; + const tsk_mutation_t *mutation; + tsk_size_t j, k, n, offset; + int code; + + if (global) { + if (num_global_alleles > TSLS_MAX_ALLELES) { + ret = TSLS_ERR_TOO_MANY_ALLELES; + goto out; + } + for (j = 0; j < num_global_alleles; j++) { + allele_length[j] = (tsk_size_t) strlen(alleles[j]); + } + } + for (j = 0; j < self->num_sites; j++) { + n = global ? num_global_alleles : 0; + offset = sites->ancestral_state_offset[j]; + code = encode_allele(sites->ancestral_state + offset, + sites->ancestral_state_offset[j + 1] - offset, list, allele_length, &n, + !global); + if (code < 0) { + ret = code; + goto out; + } + self->ancestral_state[j] = (allele_t) code; + for (k = 0; k < ts->site_mutations_length[j]; k++) { + mutation = &ts->site_mutations[j][k]; + code = encode_allele(mutation->derived_state, + mutation->derived_state_length, list, allele_length, &n, !global); + if (code < 0) { + ret = code; + goto out; + } + self->derived_state[mutation->id] = (allele_t) code; + } + /* A site always has at least two alleles, so that a haplotype can + * differ from a site with no mutations (such as a tsinfer + * ancestor's focal sites, which are not yet in the tree sequence). */ + self->num_alleles[j] = TSK_MAX(n, 2); + } +out: + return ret; +} + int -ancestor_matcher_alloc(ancestor_matcher_t *self, - const matcher_indexes_t *matcher_indexes, const double *recombination_rate, +tsls_hmm_alloc(tsls_hmm_t *self, const tsk_treeseq_t *tree_sequence, + const char **alleles, tsk_size_t num_alleles, const double *recombination_rate, const double *mismatch_rate, double likelihood_threshold, int flags) { int ret = 0; - const size_t N = matcher_indexes->num_nodes; + const size_t N = tsk_treeseq_get_num_nodes(tree_sequence); + const size_t num_mutations = tsk_treeseq_get_num_mutations(tree_sequence); - memset(self, 0, sizeof(ancestor_matcher_t)); + memset(self, 0, sizeof(tsls_hmm_t)); self->flags = flags; self->likelihood_threshold = likelihood_threshold; - self->matcher_indexes = matcher_indexes; - self->num_sites = matcher_indexes->num_sites; + self->tree_sequence = tree_sequence; + self->num_sites = tsk_treeseq_get_num_sites(tree_sequence); self->num_nodes = N; + self->sites_position = tree_sequence->tables->sites.position; + self->sequence_length = tsk_treeseq_get_sequence_length(tree_sequence); + self->num_alleles = malloc((self->num_sites + 1) * sizeof(*self->num_alleles)); + self->ancestral_state + = malloc((self->num_sites + 1) * sizeof(*self->ancestral_state)); + self->derived_state = malloc((num_mutations + 1) * sizeof(*self->derived_state)); self->recombination_rate = malloc((self->num_sites + 1) * sizeof(*self->recombination_rate)); self->mismatch_rate = malloc((self->num_sites + 1) * sizeof(*self->mismatch_rate)); @@ -583,9 +454,11 @@ ancestor_matcher_alloc(ancestor_matcher_t *self, self->likelihood_nodes_tmp = malloc((N + 1) * sizeof(*self->likelihood_nodes_tmp)); self->values_tmp = malloc((N + 1) * sizeof(*self->values_tmp)); - if (self->recombination_rate == NULL || self->mismatch_rate == NULL - || self->traceback == NULL || self->max_likelihood_node == NULL - || self->parent == NULL || self->left_child == NULL || self->right_child == NULL + if (self->num_alleles == NULL || self->ancestral_state == NULL + || self->derived_state == NULL || self->recombination_rate == NULL + || self->mismatch_rate == NULL || self->traceback == NULL + || self->max_likelihood_node == NULL || self->parent == NULL + || self->left_child == NULL || self->right_child == NULL || 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 @@ -595,6 +468,10 @@ ancestor_matcher_alloc(ancestor_matcher_t *self, ret = TSLS_ERR_NO_MEMORY; goto out; } + ret = tsls_hmm_encode_alleles(self, alleles, num_alleles); + if (ret != 0) { + goto out; + } /* Alloc in blocks of at least 64MiB, and large enough to hold the * traceback for a site in which every node has a likelihood. */ self->traceback_block_size @@ -618,8 +495,11 @@ ancestor_matcher_alloc(ancestor_matcher_t *self, } int -ancestor_matcher_free(ancestor_matcher_t *self) +tsls_hmm_free(tsls_hmm_t *self) { + tsi_safe_free(self->num_alleles); + tsi_safe_free(self->ancestral_state); + tsi_safe_free(self->derived_state); tsi_safe_free(self->recombination_rate); tsi_safe_free(self->mismatch_rate); tsi_safe_free(self->parent); @@ -644,7 +524,7 @@ ancestor_matcher_free(ancestor_matcher_t *self) } int -ancestor_matcher_set_pruning(ancestor_matcher_t *self, int mode, double parameter) +tsls_hmm_set_pruning(tsls_hmm_t *self, int mode, double parameter) { int ret = 0; @@ -680,7 +560,7 @@ cmp_double(const void *a, const void *b) /* Copy the current likelihood values into values_tmp, sorted. */ static void -ancestor_matcher_sort_values(ancestor_matcher_t *self) +tsls_hmm_sort_values(tsls_hmm_t *self) { int j; @@ -694,7 +574,7 @@ ancestor_matcher_sort_values(ancestor_matcher_t *self) /* Apply the experimental pruning, giving pruned nodes the value * likelihood_threshold. The values must be normalised. */ static void -ancestor_matcher_prune(ancestor_matcher_t *self) +tsls_hmm_prune(tsls_hmm_t *self) { double *restrict L = self->likelihood; const tsk_id_t *restrict L_nodes = self->likelihood_nodes; @@ -716,7 +596,7 @@ ancestor_matcher_prune(ancestor_matcher_t *self) if (num_likelihood_nodes <= n) { return; } - ancestor_matcher_sort_values(self); + tsls_hmm_sort_values(self); /* Keep values >= the nth largest (ties at the cutoff are kept) */ cutoff = self->values_tmp[num_likelihood_nodes - n]; } @@ -729,7 +609,7 @@ ancestor_matcher_prune(ancestor_matcher_t *self) /* Count the number of distinct likelihood values at this site. */ static void -ancestor_matcher_count_distinct(ancestor_matcher_t *self) +tsls_hmm_count_distinct(tsls_hmm_t *self) { int j; size_t num_distinct; @@ -737,7 +617,7 @@ ancestor_matcher_count_distinct(ancestor_matcher_t *self) if (!(self->flags & TSLS_COUNT_DISTINCT) || self->num_likelihood_nodes == 0) { return; } - ancestor_matcher_sort_values(self); + tsls_hmm_sort_values(self); num_distinct = 1; for (j = 1; j < self->num_likelihood_nodes; j++) { if (self->values_tmp[j] != self->values_tmp[j - 1]) { @@ -748,7 +628,7 @@ ancestor_matcher_count_distinct(ancestor_matcher_t *self) } static inline void -ancestor_matcher_add_likelihood(ancestor_matcher_t *self, tsk_id_t u, double value) +tsls_hmm_add_likelihood(tsls_hmm_t *self, tsk_id_t u, double value) { self->likelihood[u] = value; self->likelihood_nodes[self->num_likelihood_nodes] = u; @@ -757,7 +637,7 @@ ancestor_matcher_add_likelihood(ancestor_matcher_t *self, tsk_id_t u, double val /* Remove the specified node from the list of nodes with likelihoods. */ static void -ancestor_matcher_delete_likelihood(ancestor_matcher_t *self, const tsk_id_t node) +tsls_hmm_delete_likelihood(tsls_hmm_t *self, const tsk_id_t node) { int j, k; tsk_id_t *restrict L_nodes = self->likelihood_nodes; @@ -774,7 +654,7 @@ ancestor_matcher_delete_likelihood(ancestor_matcher_t *self, const tsk_id_t node } static inline void -ancestor_matcher_update_num_below(ancestor_matcher_t *self, tsk_id_t u, int64_t delta) +tsls_hmm_update_num_below(tsls_hmm_t *self, tsk_id_t u, int64_t delta) { const tsk_id_t *restrict parent = self->parent; int64_t *restrict num_below = self->num_below; @@ -787,7 +667,7 @@ ancestor_matcher_update_num_below(ancestor_matcher_t *self, tsk_id_t u, int64_t /* Returns the likelihood inherited by node u, filling in the likelihood * cache along the path so that later queries are fast. */ static double -ancestor_matcher_get_inherited_likelihood(ancestor_matcher_t *self, tsk_id_t u) +tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u) { const double *restrict L = self->likelihood; double *restrict L_cache = self->likelihood_cache; @@ -814,7 +694,7 @@ ancestor_matcher_get_inherited_likelihood(ancestor_matcher_t *self, tsk_id_t u) } static void -ancestor_matcher_reset_likelihood_cache(ancestor_matcher_t *self, tsk_id_t u) +tsls_hmm_reset_likelihood_cache(tsls_hmm_t *self, tsk_id_t u) { double *restrict L_cache = self->likelihood_cache; const tsk_id_t *restrict parent = self->parent; @@ -825,36 +705,33 @@ ancestor_matcher_reset_likelihood_cache(ancestor_matcher_t *self, tsk_id_t u) } } -/* Remove the specified edge from the tree, and if update_likelihoods is +/* Remove the edge from p to c from the tree, and if update_likelihoods is * true make sure that the child keeps its likelihood. Nodes that are * isolated after removing the edge are added to the isolated list; they * only leave the tree if they are still isolated after all the edges at - * this position have been inserted (see ancestor_matcher_next_tree). */ + * this position have been inserted (see tsls_hmm_update_tree). */ static void -ancestor_matcher_remove_edge(ancestor_matcher_t *self, edge_t edge, +tsls_hmm_remove_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, bool update_likelihoods, tsk_id_t *restrict isolated, int *num_isolated) { double *restrict L = self->likelihood; - const tsk_id_t p = edge.parent; - const tsk_id_t c = edge.child; if (self->track_num_below) { - ancestor_matcher_update_num_below(self, p, -self->num_below[c]); + tsls_hmm_update_num_below(self, p, -self->num_below[c]); } if (update_likelihoods && L[c] == NULL_LIKELIHOOD) { - ancestor_matcher_add_likelihood( - self, c, ancestor_matcher_get_inherited_likelihood(self, p)); + tsls_hmm_add_likelihood(self, c, tsls_hmm_get_inherited_likelihood(self, p)); } - remove_edge(edge, self->parent, self->left_child, self->right_child, self->left_sib, - self->right_sib); + remove_edge(p, c, self->parent, self->left_child, self->right_child, + self->left_sib, self->right_sib); if (update_likelihoods) { /* Each node can only become isolated once when removing edges, so * the isolated list has at most num_nodes entries. */ - if (!ancestor_matcher_in_tree(self, c)) { + if (!tsls_hmm_in_tree(self, c)) { isolated[*num_isolated] = c; (*num_isolated)++; } - if (!ancestor_matcher_in_tree(self, p)) { + if (!tsls_hmm_in_tree(self, p)) { /* p was a root, and so has a likelihood */ assert(L[p] >= 0); isolated[*num_isolated] = p; @@ -863,220 +740,81 @@ ancestor_matcher_remove_edge(ancestor_matcher_t *self, edge_t edge, } } -/* Insert the specified edge into the tree. Nodes entering the tree get +/* Insert the edge from p to c into the tree. Nodes entering the tree get * likelihood 0, and so can only be reached by recombination. */ static void -ancestor_matcher_insert_edge( - ancestor_matcher_t *self, edge_t edge, bool update_likelihoods) +tsls_hmm_insert_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, bool update_likelihoods) { double *restrict L = self->likelihood; if (update_likelihoods) { - if (L[edge.parent] == NOT_IN_TREE) { - ancestor_matcher_add_likelihood(self, edge.parent, 0); + if (L[p] == NOT_IN_TREE) { + tsls_hmm_add_likelihood(self, p, 0); } - if (L[edge.child] == NOT_IN_TREE) { - ancestor_matcher_add_likelihood(self, edge.child, 0); + if (L[c] == NOT_IN_TREE) { + tsls_hmm_add_likelihood(self, c, 0); } /* Otherwise the child was a root and keeps its likelihood */ - assert(L[edge.child] >= 0); + assert(L[c] >= 0); } - insert_edge(edge, self->parent, self->left_child, self->right_child, self->left_sib, - self->right_sib); + insert_edge(p, c, self->parent, self->left_child, self->right_child, + self->left_sib, self->right_sib); if (self->track_num_below) { - ancestor_matcher_update_num_below(self, edge.parent, self->num_below[edge.child]); + tsls_hmm_update_num_below(self, p, self->num_below[c]); } } -/* The current position in the sequence of trees, moving left-to-right - * (direction = 1) or right-to-left (direction = -1). */ -typedef struct { - int direction; - int_fast32_t in_index; - int_fast32_t out_index; - coordinate_t left; - coordinate_t right; -} tree_position_t; - +/* Apply the edge removals and insertions recorded in the tree position, + * which has just moved to the next tree in its direction of travel. */ static void -tree_position_init( - tree_position_t *self, const matcher_indexes_t *matcher_indexes, int direction) +tsls_hmm_update_tree(tsls_hmm_t *self, const tsk_tree_position_t *tp) { - const int_fast32_t M = (int_fast32_t) matcher_indexes->num_edges; - const coordinate_t L = matcher_indexes->sites.position[matcher_indexes->num_sites]; - - self->direction = direction; - if (direction > 0) { - self->in_index = 0; - self->out_index = 0; - self->left = 0; - self->right = 0; - } else { - self->in_index = M - 1; - self->out_index = M - 1; - self->left = L; - self->right = L; - } -} - -/* Move to the next tree in the direction of travel. */ -static void -ancestor_matcher_next_tree( - ancestor_matcher_t *self, tree_position_t *tp, bool update_likelihoods) -{ - 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; + const tsk_id_t *restrict edge_parent = self->tree_sequence->tables->edges.parent; + const tsk_id_t *restrict edge_child = self->tree_sequence->tables->edges.child; + const int direction = tp->direction; tsk_id_t *restrict isolated = self->likelihood_nodes_tmp; double *restrict L = self->likelihood; - int_fast32_t j, removal_start; + tsk_id_t j, e, u; int k, num_isolated = 0; - coordinate_t pos; - tsk_id_t u; - if (tp->direction > 0) { - pos = tp->right; - removal_start = tp->out_index; - while (tp->out_index < M && right_index[tp->out_index].right == pos) { - ancestor_matcher_remove_edge(self, right_index[tp->out_index], - update_likelihoods, isolated, &num_isolated); - tp->out_index++; - } - if (update_likelihoods) { - for (j = removal_start; j < tp->out_index; j++) { - ancestor_matcher_reset_likelihood_cache(self, right_index[j].parent); - } - } - while (tp->in_index < M && left_index[tp->in_index].left == pos) { - ancestor_matcher_insert_edge( - self, left_index[tp->in_index], update_likelihoods); - tp->in_index++; - } - tp->left = pos; - tp->right = sequence_length; - if (tp->in_index < M) { - tp->right = TSK_MIN(tp->right, left_index[tp->in_index].left); - } - if (tp->out_index < M) { - tp->right = TSK_MIN(tp->right, right_index[tp->out_index].right); - } - } else { - pos = tp->left; - removal_start = tp->out_index; - while (tp->out_index >= 0 && left_index[tp->out_index].left == pos) { - ancestor_matcher_remove_edge(self, left_index[tp->out_index], - update_likelihoods, isolated, &num_isolated); - tp->out_index--; - } - if (update_likelihoods) { - for (j = removal_start; j > tp->out_index; j--) { - ancestor_matcher_reset_likelihood_cache(self, left_index[j].parent); - } - } - while (tp->in_index >= 0 && right_index[tp->in_index].right == pos) { - ancestor_matcher_insert_edge( - self, right_index[tp->in_index], update_likelihoods); - tp->in_index--; - } - tp->right = pos; - tp->left = 0; - if (tp->out_index >= 0) { - tp->left = TSK_MAX(tp->left, left_index[tp->out_index].left); - } - if (tp->in_index >= 0) { - tp->left = TSK_MAX(tp->left, right_index[tp->in_index].right); - } + for (j = tp->out.start; j != tp->out.stop; j += direction) { + e = tp->out.order[j]; + tsls_hmm_remove_edge( + self, edge_parent[e], edge_child[e], true, isolated, &num_isolated); + } + for (j = tp->out.start; j != tp->out.stop; j += direction) { + tsls_hmm_reset_likelihood_cache(self, edge_parent[tp->out.order[j]]); + } + for (j = tp->in.start; j != tp->in.stop; j += direction) { + e = tp->in.order[j]; + tsls_hmm_insert_edge(self, edge_parent[e], edge_child[e], true); } /* Nodes that are still isolated have left the tree. Nodes that were * isolated and then inserted again keep their likelihoods. */ for (k = 0; k < num_isolated; k++) { u = isolated[k]; - if (!ancestor_matcher_in_tree(self, u)) { - ancestor_matcher_delete_likelihood(self, u); + if (!tsls_hmm_in_tree(self, u)) { + tsls_hmm_delete_likelihood(self, u); L[u] = NOT_IN_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 -ancestor_matcher_init_sweep(ancestor_matcher_t *self, tree_position_t *tp, - int direction, coordinate_t position, double initial_value) +/* Reset the HMM and build the tree containing the specified position, + * giving each root the specified initial likelihood. The tree position must + * be null. */ +static int WARN_UNUSED +tsls_hmm_init_sweep( + tsls_hmm_t *self, tsk_tree_position_t *tp, double position, double initial_value) { + int ret = 0; const size_t N = self->num_nodes; - tsk_id_t u; - size_t j; + const tsk_id_t *restrict edge_parent = self->tree_sequence->tables->edges.parent; + const tsk_id_t *restrict edge_child = self->tree_sequence->tables->edges.child; + const double *restrict edge_left = self->tree_sequence->tables->edges.left; + const double *restrict edge_right = self->tree_sequence->tables->edges.right; + tsk_id_t u, j, e; + size_t k; memset(self->parent, 0xff, N * sizeof(*self->parent)); memset(self->left_child, 0xff, N * sizeof(*self->left_child)); @@ -1084,128 +822,89 @@ ancestor_matcher_init_sweep(ancestor_matcher_t *self, tree_position_t *tp, memset(self->left_sib, 0xff, N * sizeof(*self->left_sib)); 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->likelihood_cache[j] = CACHE_UNSET; + for (k = 0; k < N; k++) { + self->likelihood_cache[k] = CACHE_UNSET; + self->num_below[k] = 1; } - tree_position_init(tp, self->matcher_indexes, direction); - if (self->track_num_below) { - for (j = 0; j < N; j++) { - self->num_below[j] = 1; + ret = seek_tree_position(tp, position); + if (ret != 0) { + goto out; + } + for (j = tp->in.start; j != tp->in.stop; j += tp->direction) { + e = tp->in.order[j]; + if (edge_left[e] <= position && position < edge_right[e]) { + tsls_hmm_insert_edge(self, edge_parent[e], edge_child[e], false); } - 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++) { if (self->parent[u] != NULL_NODE) { self->likelihood[u] = NULL_LIKELIHOOD; } else if (self->left_child[u] != NULL_NODE) { - ancestor_matcher_add_likelihood(self, u, initial_value); + tsls_hmm_add_likelihood(self, u, initial_value); } else { self->likelihood[u] = NOT_IN_TREE; } } +out: + return ret; } static inline void -ancestor_matcher_set_allelic_state(ancestor_matcher_t *self, const tsk_id_t site) -{ - mutation_list_node_t *mutation; - - for (mutation = self->matcher_indexes->sites.mutations[site]; mutation != NULL; - mutation = mutation->next) { - self->allelic_state[mutation->node] = mutation->derived_state; - } -} - -static inline void -ancestor_matcher_unset_allelic_state(ancestor_matcher_t *self, const tsk_id_t site) -{ - mutation_list_node_t *mutation; - - for (mutation = self->matcher_indexes->sites.mutations[site]; mutation != NULL; - mutation = mutation->next) { - self->allelic_state[mutation->node] = TSK_NULL; - } -} - -/* 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) +tsls_hmm_set_allelic_state(tsls_hmm_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; + const tsk_mutation_t *restrict mutations = self->tree_sequence->site_mutations[site]; + const tsk_size_t num_mutations = self->tree_sequence->site_mutations_length[site]; + tsk_size_t j; - for (j = 0; j < self->num_likelihood_nodes; j++) { - u = L_nodes[j]; - if (parent[u] == NULL_NODE) { - allelic_state[u] = ancestral_state; - } + for (j = 0; j < num_mutations; j++) { + self->allelic_state[mutations[j].node] = self->derived_state[mutations[j].id]; } } static inline void -ancestor_matcher_unset_root_allelic_state(ancestor_matcher_t *self) +tsls_hmm_unset_allelic_state(tsls_hmm_t *self, const tsk_id_t 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; + const tsk_mutation_t *restrict mutations = self->tree_sequence->site_mutations[site]; + const tsk_size_t num_mutations = self->tree_sequence->site_mutations_length[site]; + tsk_size_t j; - for (j = 0; j < self->num_likelihood_nodes; j++) { - u = L_nodes[j]; - if (parent[u] == NULL_NODE) { - allelic_state[u] = TSK_NULL; - } + for (j = 0; j < num_mutations; j++) { + self->allelic_state[mutations[j].node] = TSK_NULL; } } /* Make sure that every node with a mutation at this site that is in the * tree has its own likelihood. */ static void -ancestor_matcher_insert_mutation_likelihoods( - ancestor_matcher_t *self, const tsk_id_t site) +tsls_hmm_insert_mutation_likelihoods(tsls_hmm_t *self, const tsk_id_t site) { const double *restrict L = self->likelihood; const tsk_id_t *restrict parent = self->parent; - mutation_list_node_t *mutation; - tsk_id_t u; + const tsk_mutation_t *restrict mutations = self->tree_sequence->site_mutations[site]; + const tsk_size_t num_mutations = self->tree_sequence->site_mutations_length[site]; + tsk_size_t j; + tsk_id_t u, node; - for (mutation = self->matcher_indexes->sites.mutations[site]; mutation != NULL; - mutation = mutation->next) { - if (L[mutation->node] == NULL_LIKELIHOOD) { - u = mutation->node; + for (j = 0; j < num_mutations; j++) { + node = mutations[j].node; + if (L[node] == NULL_LIKELIHOOD) { + u = node; while (L[u] == NULL_LIKELIHOOD) { u = parent[u]; assert(u != NULL_NODE); } - ancestor_matcher_add_likelihood(self, mutation->node, L[u]); + tsls_hmm_add_likelihood(self, node, L[u]); } } } -/* As ancestor_matcher_get_inherited_likelihood, but also return the +/* As tsls_hmm_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) +tsls_hmm_get_inherited_likelihood_and_flag( + tsls_hmm_t *self, tsk_id_t u, int8_t *flag) { const double *restrict L = self->likelihood; const int8_t *restrict R = self->recombination_required; @@ -1245,8 +944,8 @@ ancestor_matcher_get_inherited_likelihood_and_flag( * 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, bool check_recombination_required) +tsls_hmm_coalesce_likelihoods( + tsls_hmm_t *self, bool check_recombination_required) { const int8_t *restrict R = self->recombination_required; double L_p; @@ -1270,13 +969,13 @@ ancestor_matcher_coalesce_likelihoods( 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( + L_p = tsls_hmm_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); + L_p = tsls_hmm_get_inherited_likelihood(self, p); if (L[u] == L_p) { L[u] = NULL_LIKELIHOOD; } @@ -1289,7 +988,7 @@ ancestor_matcher_coalesce_likelihoods( } self->num_likelihood_nodes = num_likelihood_nodes; for (j = 0; j < num_cached_paths; j++) { - ancestor_matcher_reset_likelihood_cache(self, cached_paths[j]); + tsls_hmm_reset_likelihood_cache(self, cached_paths[j]); } } @@ -1303,9 +1002,9 @@ emission_probability(allele_t node_state, allele_t state, double mu, double num_ } static inline int -check_haplotype_state(const ancestor_matcher_t *self, tsk_id_t site, allele_t state) +check_haplotype_state(const tsls_hmm_t *self, tsk_id_t site, allele_t state) { - const tsk_size_t num_alleles = self->matcher_indexes->sites.num_alleles[site]; + const tsk_size_t num_alleles = self->num_alleles[site]; if (state < TSLS_MISSING_DATA || (state >= 0 && (tsk_size_t) state >= num_alleles)) { return TSLS_ERR_BAD_HAPLOTYPE_ALLELE; @@ -1314,7 +1013,7 @@ check_haplotype_state(const ancestor_matcher_t *self, tsk_id_t site, allele_t st } static inline double -ancestor_matcher_get_n(const ancestor_matcher_t *self) +tsls_hmm_get_n(const tsls_hmm_t *self) { double n = 1.0; if (!(self->flags & TSLS_DISABLE_WEIGHT_BY_N)) { @@ -1326,7 +1025,7 @@ ancestor_matcher_get_n(const ancestor_matcher_t *self) /* Compute the number of nodes in the tree that inherit the likelihood of * each likelihood node. */ static void -ancestor_matcher_compute_num_inheriting(ancestor_matcher_t *self) +tsls_hmm_compute_num_inheriting(tsls_hmm_t *self) { const double *restrict L = self->likelihood; const tsk_id_t *restrict parent = self->parent; @@ -1354,7 +1053,7 @@ ancestor_matcher_compute_num_inheriting(ancestor_matcher_t *self) /* Returns the sum of the likelihoods over all the nodes in the tree. */ static double -ancestor_matcher_total_likelihood(const ancestor_matcher_t *self) +tsls_hmm_total_likelihood(const tsls_hmm_t *self) { const double *restrict L = self->likelihood; const tsk_id_t *restrict L_nodes = self->likelihood_nodes; @@ -1373,7 +1072,7 @@ ancestor_matcher_total_likelihood(const ancestor_matcher_t *self) /* Store the recombination_required state in the traceback */ static int WARN_UNUSED -ancestor_matcher_store_traceback(ancestor_matcher_t *self, const tsk_id_t site_id) +tsls_hmm_store_traceback(tsls_hmm_t *self, const tsk_id_t site_id) { int ret = 0; tsk_id_t u; @@ -1433,7 +1132,7 @@ ancestor_matcher_store_traceback(ancestor_matcher_t *self, const tsk_id_t site_i } static int WARN_UNUSED -ancestor_matcher_process_site_viterbi(ancestor_matcher_t *self, const tsk_id_t site, +tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site, const allele_t state, compressed_matrix_t *output) { int ret = 0; @@ -1445,15 +1144,15 @@ ancestor_matcher_process_site_viterbi(ancestor_matcher_t *self, const tsk_id_t s int8_t *restrict recombination_required = self->recombination_required; const double rho = self->recombination_rate[site]; const double mu = self->mismatch_rate[site]; - const double n = ancestor_matcher_get_n(self); - const double num_alleles = (double) self->matcher_indexes->sites.num_alleles[site]; + const double n = tsls_hmm_get_n(self); + const double num_alleles = (double) self->num_alleles[site]; + const allele_t ancestral_state = self->ancestral_state[site]; double max_L, p_no_recomb, p_recomb, p_t, p_e; tsk_id_t u, max_L_node; int j; - ancestor_matcher_insert_mutation_likelihoods(self, site); - ancestor_matcher_set_root_allelic_state(self, site); - ancestor_matcher_set_allelic_state(self, site); + tsls_hmm_insert_mutation_likelihoods(self, site); + tsls_hmm_set_allelic_state(self, site); max_L = -1; max_L_node = NULL_NODE; for (j = 0; j < self->num_likelihood_nodes; j++) { @@ -1468,15 +1167,15 @@ ancestor_matcher_process_site_viterbi(ancestor_matcher_t *self, const tsk_id_t s recombination_required[u] = true; } p_e = emission_probability( - get_marked_allelic_state(allelic_state, parent, u), state, mu, num_alleles); + get_allelic_state(allelic_state, parent, u, ancestral_state), state, mu, + num_alleles); L[u] = p_t * p_e; if (L[u] > max_L) { max_L = L[u]; max_L_node = u; } } - ancestor_matcher_unset_allelic_state(self, site); - ancestor_matcher_unset_root_allelic_state(self); + tsls_hmm_unset_allelic_state(self, site); (void) num_likelihood_nodes; if (max_L <= 0) { @@ -1495,13 +1194,13 @@ ancestor_matcher_process_site_viterbi(ancestor_matcher_t *self, const tsk_id_t s u = L_nodes[j]; L[u] = TSK_MAX(L[u] / max_L, self->likelihood_threshold); } - ancestor_matcher_prune(self); - ancestor_matcher_count_distinct(self); + tsls_hmm_prune(self); + tsls_hmm_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); + tsls_hmm_coalesce_likelihoods(self, true); + ret = tsls_hmm_store_traceback(self, site); if (ret != 0) { goto out; } @@ -1521,7 +1220,7 @@ ancestor_matcher_process_site_viterbi(ancestor_matcher_t *self, const tsk_id_t s * compute the total at the end of each site and use it at the next. At the * first site, previous_total is negative and we use the current total. */ static int WARN_UNUSED -ancestor_matcher_process_site_forward(ancestor_matcher_t *self, const tsk_id_t site, +tsls_hmm_process_site_forward(tsls_hmm_t *self, const tsk_id_t site, const allele_t state, double *previous_total, compressed_matrix_t *output) { int ret = 0; @@ -1531,29 +1230,29 @@ ancestor_matcher_process_site_forward(ancestor_matcher_t *self, const tsk_id_t s double *restrict L = self->likelihood; const double rho = self->recombination_rate[site]; const double mu = self->mismatch_rate[site]; - const double n = ancestor_matcher_get_n(self); - const double num_alleles = (double) self->matcher_indexes->sites.num_alleles[site]; + const double n = tsls_hmm_get_n(self); + const double num_alleles = (double) self->num_alleles[site]; + const allele_t ancestral_state = self->ancestral_state[site]; double total, normalisation_factor, p_e; tsk_id_t u; int j; - ancestor_matcher_insert_mutation_likelihoods(self, site); - ancestor_matcher_compute_num_inheriting(self); + tsls_hmm_insert_mutation_likelihoods(self, site); + tsls_hmm_compute_num_inheriting(self); total = *previous_total; if (total < 0) { - total = ancestor_matcher_total_likelihood(self); + total = tsls_hmm_total_likelihood(self); } - ancestor_matcher_set_root_allelic_state(self, site); - ancestor_matcher_set_allelic_state(self, site); + tsls_hmm_set_allelic_state(self, site); for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; p_e = emission_probability( - get_marked_allelic_state(allelic_state, parent, u), state, mu, num_alleles); + get_allelic_state(allelic_state, parent, u, ancestral_state), 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); + tsls_hmm_unset_allelic_state(self, site); + normalisation_factor = tsls_hmm_total_likelihood(self); if (!(normalisation_factor > 0)) { ret = TSLS_ERR_MATCH_IMPOSSIBLE; goto out; @@ -1562,14 +1261,14 @@ ancestor_matcher_process_site_forward(ancestor_matcher_t *self, const tsk_id_t s u = L_nodes[j]; L[u] = TSK_MAX(L[u] / normalisation_factor, self->likelihood_threshold); } - ancestor_matcher_prune(self); - *previous_total = ancestor_matcher_total_likelihood(self); + tsls_hmm_prune(self); + *previous_total = tsls_hmm_total_likelihood(self); ret = compressed_matrix_store_site(output, site, normalisation_factor, self->num_likelihood_nodes, L_nodes, L, NULL); if (ret != 0) { goto out; } - ancestor_matcher_coalesce_likelihoods(self, false); + tsls_hmm_coalesce_likelihoods(self, false); out: return ret; } @@ -1586,7 +1285,7 @@ typedef struct { } backward_transition_t; static int WARN_UNUSED -ancestor_matcher_process_site_backward(ancestor_matcher_t *self, const tsk_id_t site, +tsls_hmm_process_site_backward(tsls_hmm_t *self, const tsk_id_t site, const allele_t state, const double *forward_normalisation_factor, backward_transition_t *transition, compressed_matrix_t *output) { @@ -1596,8 +1295,9 @@ ancestor_matcher_process_site_backward(ancestor_matcher_t *self, const tsk_id_t const allele_t *restrict allelic_state = self->allelic_state; double *restrict L = self->likelihood; const double mu = self->mismatch_rate[site]; - const double n = ancestor_matcher_get_n(self); - const double num_alleles = (double) self->matcher_indexes->sites.num_alleles[site]; + const double n = tsls_hmm_get_n(self); + const double num_alleles = (double) self->num_alleles[site]; + const allele_t ancestral_state = self->ancestral_state[site]; double x; tsk_id_t u; int j; @@ -1609,47 +1309,48 @@ ancestor_matcher_process_site_backward(ancestor_matcher_t *self, const tsk_id_t L[u] = TSK_MAX( x / transition->normalisation_factor, self->likelihood_threshold); } - ancestor_matcher_prune(self); + tsls_hmm_prune(self); } - ancestor_matcher_insert_mutation_likelihoods(self, site); + tsls_hmm_insert_mutation_likelihoods(self, site); ret = compressed_matrix_store_site(output, site, forward_normalisation_factor[site], self->num_likelihood_nodes, L_nodes, L, NULL); if (ret != 0) { goto out; } - ancestor_matcher_compute_num_inheriting(self); - ancestor_matcher_set_root_allelic_state(self, site); - ancestor_matcher_set_allelic_state(self, site); + tsls_hmm_compute_num_inheriting(self); + tsls_hmm_set_allelic_state(self, site); for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; L[u] *= emission_probability( - get_marked_allelic_state(allelic_state, parent, u), state, mu, num_alleles); + get_allelic_state(allelic_state, parent, u, ancestral_state), state, mu, + num_alleles); } - ancestor_matcher_unset_allelic_state(self, site); - ancestor_matcher_unset_root_allelic_state(self); + tsls_hmm_unset_allelic_state(self, site); transition->pending = true; - transition->total = ancestor_matcher_total_likelihood(self); + transition->total = tsls_hmm_total_likelihood(self); transition->rho = self->recombination_rate[site]; transition->normalisation_factor = forward_normalisation_factor[site]; - ancestor_matcher_coalesce_likelihoods(self, false); + tsls_hmm_coalesce_likelihoods(self, false); out: return ret; } /* Run the specified algorithm over sites [start, end). */ static int WARN_UNUSED -ancestor_matcher_run(ancestor_matcher_t *self, int mode, tsk_id_t start, tsk_id_t end, +tsls_hmm_run(tsls_hmm_t *self, int mode, tsk_id_t start, tsk_id_t end, const allele_t *haplotype, const double *forward_normalisation_factor, compressed_matrix_t *output) { int ret = 0; - tree_position_t tp; + tsk_tree_position_t tp; backward_transition_t transition; - const coordinate_t *restrict position = self->matcher_indexes->sites.position; - const int direction = mode == MODE_BACKWARD ? -1 : 1; + const double *restrict position = self->sites_position; + const int direction = mode == MODE_BACKWARD ? TSK_DIR_REVERSE : TSK_DIR_FORWARD; tsk_id_t site, first_site; double previous_total = -1; + bool valid; + tsk_tree_position_init(&tp, self->tree_sequence, 0); for (site = start; site < end; site++) { ret = check_haplotype_state(self, site, haplotype[site]); if (ret != 0) { @@ -1658,24 +1359,31 @@ ancestor_matcher_run(ancestor_matcher_t *self, int mode, tsk_id_t start, tsk_id_ } self->track_num_below = mode != MODE_VITERBI; first_site = direction > 0 ? start : end - 1; - ancestor_matcher_init_sweep(self, &tp, direction, position[first_site], 1.0); + ret = tsls_hmm_init_sweep(self, &tp, position[first_site], 1.0); + if (ret != 0) { + goto out; + } transition.pending = false; for (site = first_site; start <= site && site < end; site += direction) { - while (!(tp.left <= position[site] && position[site] < tp.right)) { - ancestor_matcher_next_tree(self, &tp, true); + while (!(tp.interval.left <= position[site] + && position[site] < tp.interval.right)) { + valid = direction > 0 ? tsk_tree_position_next(&tp) + : tsk_tree_position_prev(&tp); + assert(valid); + (void) valid; + tsls_hmm_update_tree(self, &tp); } if (self->flags & TSLS_EXTENDED_CHECKS) { - ancestor_matcher_check_state(self); + tsls_hmm_check_state(self); } if (mode == MODE_VITERBI) { - ret = ancestor_matcher_process_site_viterbi( - self, site, haplotype[site], output); + ret = tsls_hmm_process_site_viterbi(self, site, haplotype[site], output); } else if (mode == MODE_FORWARD) { - ret = ancestor_matcher_process_site_forward( + ret = tsls_hmm_process_site_forward( self, site, haplotype[site], &previous_total, output); } else { - ret = ancestor_matcher_process_site_backward(self, site, haplotype[site], + ret = tsls_hmm_process_site_backward(self, site, haplotype[site], forward_normalisation_factor, &transition, output); } if (ret != 0) { @@ -1683,14 +1391,15 @@ ancestor_matcher_run(ancestor_matcher_t *self, int mode, tsk_id_t start, tsk_id_ } } out: + tsk_tree_position_free(&tp); return ret; } /* Resets the recombination_required array from the traceback at the specified site. */ static inline void -ancestor_matcher_set_recombination_required( - ancestor_matcher_t *self, tsk_id_t site, int8_t *restrict recombination_required) +tsls_hmm_set_recombination_required( + tsls_hmm_t *self, tsk_id_t site, int8_t *restrict recombination_required) { int j; const int8_t *restrict R = self->traceback[site].recombination_required; @@ -1705,8 +1414,8 @@ ancestor_matcher_set_recombination_required( /* Unsets the likelihood array from the traceback at the specified site. */ static inline void -ancestor_matcher_unset_recombination_required( - ancestor_matcher_t *self, tsk_id_t site, int8_t *restrict recombination_required) +tsls_hmm_unset_recombination_required( + tsls_hmm_t *self, tsk_id_t site, int8_t *restrict recombination_required) { int j; const tsk_id_t *restrict node = self->traceback[site].node; @@ -1717,30 +1426,33 @@ ancestor_matcher_unset_recombination_required( } } -static void -ancestor_matcher_run_traceback(ancestor_matcher_t *self, tsk_id_t start, tsk_id_t end, - allele_t *restrict match, size_t *path_length_out, tsk_id_t *restrict path_left, - tsk_id_t *restrict path_right, tsk_id_t *restrict path_parent) +/* Run the traceback over sites [start, end), moving right-to-left through + * the trees. Only the parent array of the tree is needed. */ +static int WARN_UNUSED +tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, + allele_t *restrict match, size_t *path_length_out, double *restrict path_left, + double *restrict path_right, tsk_id_t *restrict path_parent) { - tsk_id_t site; - edge_t edge; - tsk_id_t u, max_likelihood_node; - coordinate_t left, right, pos, start_pos, end_pos; + int ret = 0; + tsk_tree_position_t tp; + tsk_id_t site, j, e, u, max_likelihood_node; + double x; + bool valid; tsk_id_t *restrict parent = self->parent; const allele_t *restrict allelic_state = self->allelic_state; int8_t *restrict recombination_required = self->recombination_required; - const edge_t *restrict in = self->matcher_indexes->right_index_edges; - const edge_t *restrict out = self->matcher_indexes->left_index_edges; - const coordinate_t *restrict sites_position = self->matcher_indexes->sites.position; - const allele_t *restrict ancestral_state - = self->matcher_indexes->sites.ancestral_state; - const coordinate_t sequence_length = sites_position[self->num_sites]; - int_fast32_t in_index = (int_fast32_t) self->matcher_indexes->num_edges - 1; - int_fast32_t out_index = (int_fast32_t) self->matcher_indexes->num_edges - 1; + const double *restrict sites_position = self->sites_position; + const tsk_edge_table_t *edges = &self->tree_sequence->tables->edges; + const tsk_id_t *restrict edge_parent = edges->parent; + const tsk_id_t *restrict edge_child = edges->child; + const double *restrict edge_left = edges->left; + const double *restrict edge_right = edges->right; + const double start_pos = start == 0 ? 0 : sites_position[start]; + const double end_pos + = end == (tsk_id_t) self->num_sites ? self->sequence_length : sites_position[end]; size_t path_length = 0; - start_pos = start == 0 ? 0 : sites_position[start]; - end_pos = sites_position[end]; + tsk_tree_position_init(&tp, self->tree_sequence, 0); /* Prepare for the traceback and get the memory ready for recording * the output edges. */ @@ -1749,78 +1461,73 @@ ancestor_matcher_run_traceback(ancestor_matcher_t *self, tsk_id_t start, tsk_id_ assert(max_likelihood_node != NULL_NODE); path_parent[path_length] = max_likelihood_node; - /* Now go through the trees in reverse and run the traceback */ + /* Build the tree at the last site, and then go through the trees in + * reverse, running the traceback at each site. */ memset(parent, 0xff, self->num_nodes * sizeof(*parent)); memset( recombination_required, 0xff, self->num_nodes * sizeof(*recombination_required)); - pos = sequence_length; - site = (tsk_id_t) self->num_sites - 1; - - while (pos > start_pos) { - while (out_index >= 0 && out[out_index].left == pos) { - edge = out[out_index]; - out_index--; - parent[edge.child] = NULL_NODE; - } - while (in_index >= 0 && in[in_index].right == pos) { - edge = in[in_index]; - in_index--; - parent[edge.child] = edge.parent; - } - right = pos; - left = 0; - if (out_index >= 0) { - left = TSK_MAX(left, out[out_index].left); - } - if (in_index >= 0) { - left = TSK_MAX(left, in[in_index].right); + x = sites_position[end - 1]; + ret = seek_tree_position(&tp, x); + if (ret != 0) { + goto out; + } + for (j = tp.in.start; j != tp.in.stop; j += tp.direction) { + e = tp.in.order[j]; + if (edge_left[e] <= x && x < edge_right[e]) { + parent[edge_child[e]] = edge_parent[e]; } - pos = left; - - /* The tree is ready; perform the traceback at each site in this tree */ - assert(left < right); - for (; site >= 0 && left <= sites_position[site] && sites_position[site] < right; - site--) { - if (start_pos <= sites_position[site] && sites_position[site] < end_pos) { - ancestor_matcher_set_allelic_state(self, site); - u = path_parent[path_length]; - match[site] = get_allelic_state( - allelic_state, parent, u, ancestral_state[site]); - ancestor_matcher_unset_allelic_state(self, site); - - /* Mark the traceback nodes on the tree */ - ancestor_matcher_set_recombination_required( - self, site, recombination_required); - - /* Traverse up the tree from the current node. The first marked - * node that we meet tells us whether we need to recombine. Every - * root has a likelihood, and so is marked. */ - while (recombination_required[u] == -1) { - u = parent[u]; - assert(u != NULL_NODE); - } - if (recombination_required[u] && site > start) { - max_likelihood_node = self->max_likelihood_node[site - 1]; - assert(max_likelihood_node != NULL_NODE); - path_left[path_length] = sites_position[site]; - path_length++; - /* Start the next output edge */ - path_right[path_length] = sites_position[site]; - path_parent[path_length] = max_likelihood_node; - } - /* Unset the values in the tree for the next site. */ - ancestor_matcher_unset_recombination_required( - self, site, recombination_required); + } + for (site = end - 1; site >= start; site--) { + while (sites_position[site] < tp.interval.left) { + valid = tsk_tree_position_prev(&tp); + assert(valid); + (void) valid; + for (j = tp.out.start; j != tp.out.stop; j--) { + parent[edge_child[tp.out.order[j]]] = NULL_NODE; + } + for (j = tp.in.start; j != tp.in.stop; j--) { + e = tp.in.order[j]; + parent[edge_child[e]] = edge_parent[e]; } } + tsls_hmm_set_allelic_state(self, site); + u = path_parent[path_length]; + match[site] + = get_allelic_state(allelic_state, parent, u, self->ancestral_state[site]); + tsls_hmm_unset_allelic_state(self, site); + + /* Mark the traceback nodes on the tree */ + tsls_hmm_set_recombination_required(self, site, recombination_required); + + /* Traverse up the tree from the current node. The first marked + * node that we meet tells us whether we need to recombine. Every + * root has a likelihood, and so is marked. */ + while (recombination_required[u] == -1) { + u = parent[u]; + assert(u != NULL_NODE); + } + if (recombination_required[u] && site > start) { + max_likelihood_node = self->max_likelihood_node[site - 1]; + assert(max_likelihood_node != NULL_NODE); + path_left[path_length] = sites_position[site]; + path_length++; + /* Start the next output edge */ + path_right[path_length] = sites_position[site]; + path_parent[path_length] = max_likelihood_node; + } + /* Unset the values in the tree for the next site. */ + tsls_hmm_unset_recombination_required(self, site, recombination_required); } path_left[path_length] = start_pos; path_length++; *path_length_out = path_length; +out: + tsk_tree_position_free(&tp); + return ret; } static int -ancestor_matcher_reset_traceback(ancestor_matcher_t *self) +tsls_hmm_reset_traceback(tsls_hmm_t *self) { int ret = 0; @@ -1845,9 +1552,9 @@ ancestor_matcher_reset_traceback(ancestor_matcher_t *self) * traceback at each site are also stored in output, with the maximum * likelihood (before normalisation) as the normalisation factor. */ int -ancestor_matcher_find_path_matrix(ancestor_matcher_t *self, tsk_id_t start, +tsls_hmm_find_path_matrix(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, const allele_t *haplotype, allele_t *matched_haplotype, - size_t *path_length, tsk_id_t *path_left, tsk_id_t *path_right, + size_t *path_length, double *path_left, double *path_right, tsk_id_t *path_parent, compressed_matrix_t *output) { int ret = 0; @@ -1860,15 +1567,15 @@ ancestor_matcher_find_path_matrix(ancestor_matcher_t *self, tsk_id_t start, ret = TSLS_ERR_BAD_SITE_RANGE; goto out; } - ret = ancestor_matcher_reset_traceback(self); + ret = tsls_hmm_reset_traceback(self); if (ret != 0) { goto out; } - ret = ancestor_matcher_run(self, MODE_VITERBI, start, end, haplotype, NULL, output); + ret = tsls_hmm_run(self, MODE_VITERBI, start, end, haplotype, NULL, output); if (ret != 0) { goto out; } - ancestor_matcher_run_traceback(self, start, end, matched_haplotype, path_length, + ret = tsls_hmm_run_traceback(self, start, end, matched_haplotype, path_length, path_left, path_right, path_parent); out: /* Reset some memory for the next call */ @@ -1880,17 +1587,17 @@ ancestor_matcher_find_path_matrix(ancestor_matcher_t *self, tsk_id_t start, } int -ancestor_matcher_find_path(ancestor_matcher_t *self, tsk_id_t start, tsk_id_t end, +tsls_hmm_find_path(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, const allele_t *haplotype, allele_t *matched_haplotype, size_t *path_length, - tsk_id_t *path_left, tsk_id_t *path_right, tsk_id_t *path_parent) + double *path_left, double *path_right, tsk_id_t *path_parent) { - return ancestor_matcher_find_path_matrix(self, start, end, haplotype, + return tsls_hmm_find_path_matrix(self, start, end, haplotype, matched_haplotype, path_length, path_left, path_right, path_parent, NULL); } int -ancestor_matcher_forward( - ancestor_matcher_t *self, const allele_t *haplotype, compressed_matrix_t *output) +tsls_hmm_forward( + tsls_hmm_t *self, const allele_t *haplotype, compressed_matrix_t *output) { int ret = 0; @@ -1899,7 +1606,7 @@ ancestor_matcher_forward( goto out; } if (self->num_sites > 0) { - ret = ancestor_matcher_run(self, MODE_FORWARD, 0, (tsk_id_t) self->num_sites, + ret = tsls_hmm_run(self, MODE_FORWARD, 0, (tsk_id_t) self->num_sites, haplotype, NULL, output); } out: @@ -1908,7 +1615,7 @@ ancestor_matcher_forward( } int -ancestor_matcher_backward(ancestor_matcher_t *self, const allele_t *haplotype, +tsls_hmm_backward(tsls_hmm_t *self, const allele_t *haplotype, const double *forward_normalisation_factor, compressed_matrix_t *output) { int ret = 0; @@ -1925,7 +1632,7 @@ ancestor_matcher_backward(ancestor_matcher_t *self, const allele_t *haplotype, } } if (self->num_sites > 0) { - ret = ancestor_matcher_run(self, MODE_BACKWARD, 0, (tsk_id_t) self->num_sites, + ret = tsls_hmm_run(self, MODE_BACKWARD, 0, (tsk_id_t) self->num_sites, haplotype, forward_normalisation_factor, output); } out: @@ -1934,7 +1641,7 @@ ancestor_matcher_backward(ancestor_matcher_t *self, const allele_t *haplotype, } double -ancestor_matcher_get_mean_traceback_size(ancestor_matcher_t *self) +tsls_hmm_get_mean_traceback_size(tsls_hmm_t *self) { if (self->num_sites == 0) { return 0; @@ -1943,7 +1650,7 @@ ancestor_matcher_get_mean_traceback_size(ancestor_matcher_t *self) } double -ancestor_matcher_get_mean_num_distinct_values(ancestor_matcher_t *self) +tsls_hmm_get_mean_num_distinct_values(tsls_hmm_t *self) { if (self->num_sites == 0) { return 0; @@ -1952,7 +1659,7 @@ ancestor_matcher_get_mean_num_distinct_values(ancestor_matcher_t *self) } size_t -ancestor_matcher_get_total_memory(ancestor_matcher_t *self) +tsls_hmm_get_total_memory(tsls_hmm_t *self) { return self->traceback_allocator.total_size; } diff --git a/tsls/lib/meson.build b/tsls/lib/meson.build index 7e27f16..39f182b 100644 --- a/tsls/lib/meson.build +++ b/tsls/lib/meson.build @@ -16,7 +16,7 @@ extra_c_args = [ '-Wcast-qual', '-Wwrite-strings', '-Wnested-externs', '-fshort-enums', '-fno-common'] -tsls_sources = ['matcher.c', 'err.c', 'tsi_blkalloc.c'] +tsls_sources = ['hmm.c', 'err.c', 'tsi_blkalloc.c'] tsls_lib = static_library('tsls', sources: tsls_sources, dependencies: [m_dep, tskit_dep], diff --git a/tsls/lib/tests/tests.c b/tsls/lib/tests/tests.c index 49b514d..fd40a73 100644 --- a/tsls/lib/tests/tests.c +++ b/tsls/lib/tests/tests.c @@ -13,129 +13,141 @@ static FILE *_devnull; /* A forest over [0, 10): roots 3 (above 0 and 1) and 4 (above 2) on [0, 5), - * joined by node 6 on [5, 10). Node 5 is isolated. */ + * joined by node 6 on [5, 10). Node 5 is isolated. If extra_site is true, + * add a fifth site with no mutations. */ static void -build_forest(tsk_table_collection_t *tables) +build_forest_ts(tsk_treeseq_t *ts, bool extra_site) { int ret; int j; double time[] = { 0, 0, 0, 1, 2, 4, 3 }; + tsk_table_collection_t tables; - ret = tsk_table_collection_init(tables, 0); + ret = tsk_table_collection_init(&tables, 0); CU_ASSERT_EQUAL_FATAL(ret, 0); - tables->sequence_length = 10; + tables.sequence_length = 10; for (j = 0; j < 7; j++) { - tsk_node_table_add_row(&tables->nodes, time[j] == 0 ? TSK_NODE_IS_SAMPLE : 0, + tsk_node_table_add_row(&tables.nodes, time[j] == 0 ? TSK_NODE_IS_SAMPLE : 0, time[j], TSK_NULL, TSK_NULL, NULL, 0); } - tsk_edge_table_add_row(&tables->edges, 0, 10, 3, 0, NULL, 0); - tsk_edge_table_add_row(&tables->edges, 0, 10, 3, 1, NULL, 0); - tsk_edge_table_add_row(&tables->edges, 0, 10, 4, 2, NULL, 0); - tsk_edge_table_add_row(&tables->edges, 5, 10, 6, 3, NULL, 0); - tsk_edge_table_add_row(&tables->edges, 5, 10, 6, 4, NULL, 0); - tsk_site_table_add_row(&tables->sites, 1, "A", 1, NULL, 0); - tsk_site_table_add_row(&tables->sites, 3, "C", 1, NULL, 0); - tsk_site_table_add_row(&tables->sites, 6, "G", 1, NULL, 0); - tsk_site_table_add_row(&tables->sites, 8, "A", 1, NULL, 0); + tsk_edge_table_add_row(&tables.edges, 0, 10, 3, 0, NULL, 0); + tsk_edge_table_add_row(&tables.edges, 0, 10, 3, 1, NULL, 0); + tsk_edge_table_add_row(&tables.edges, 0, 10, 4, 2, NULL, 0); + tsk_edge_table_add_row(&tables.edges, 5, 10, 6, 3, NULL, 0); + tsk_edge_table_add_row(&tables.edges, 5, 10, 6, 4, NULL, 0); + tsk_site_table_add_row(&tables.sites, 1, "A", 1, NULL, 0); + tsk_site_table_add_row(&tables.sites, 3, "C", 1, NULL, 0); + tsk_site_table_add_row(&tables.sites, 6, "G", 1, NULL, 0); + tsk_site_table_add_row(&tables.sites, 8, "A", 1, NULL, 0); tsk_mutation_table_add_row( - &tables->mutations, 0, 3, TSK_NULL, TSK_UNKNOWN_TIME, "C", 1, NULL, 0); + &tables.mutations, 0, 3, TSK_NULL, TSK_UNKNOWN_TIME, "C", 1, NULL, 0); tsk_mutation_table_add_row( - &tables->mutations, 0, 2, TSK_NULL, TSK_UNKNOWN_TIME, "G", 1, NULL, 0); + &tables.mutations, 0, 2, TSK_NULL, TSK_UNKNOWN_TIME, "G", 1, NULL, 0); tsk_mutation_table_add_row( - &tables->mutations, 1, 4, TSK_NULL, TSK_UNKNOWN_TIME, "A", 1, NULL, 0); + &tables.mutations, 1, 4, TSK_NULL, TSK_UNKNOWN_TIME, "A", 1, NULL, 0); tsk_mutation_table_add_row( - &tables->mutations, 2, 0, TSK_NULL, TSK_UNKNOWN_TIME, "A", 1, NULL, 0); + &tables.mutations, 2, 0, TSK_NULL, TSK_UNKNOWN_TIME, "A", 1, NULL, 0); tsk_mutation_table_add_row( - &tables->mutations, 3, 6, TSK_NULL, TSK_UNKNOWN_TIME, "T", 1, NULL, 0); + &tables.mutations, 3, 6, TSK_NULL, TSK_UNKNOWN_TIME, "T", 1, NULL, 0); tsk_mutation_table_add_row( - &tables->mutations, 3, 1, TSK_NULL, TSK_UNKNOWN_TIME, "A", 1, NULL, 0); - ret = tsk_table_collection_sort(tables, NULL, 0); + &tables.mutations, 3, 1, TSK_NULL, TSK_UNKNOWN_TIME, "A", 1, NULL, 0); + if (extra_site) { + tsk_site_table_add_row(&tables.sites, 9, "T", 1, NULL, 0); + } + ret = tsk_table_collection_sort(&tables, NULL, 0); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_table_collection_build_index(tables, 0); + ret = tsk_treeseq_init(ts, &tables, + TSK_TS_INIT_BUILD_INDEXES | TSK_TS_INIT_COMPUTE_MUTATION_PARENTS); CU_ASSERT_EQUAL_FATAL(ret, 0); + tsk_table_collection_free(&tables); } -/* Alleles are A=0, C=1, G=2, T=3 */ -static tsk_size_t num_alleles[] = { 4, 4, 4, 4 }; -static allele_t ancestral_state[] = { 0, 1, 2, 0 }; -static allele_t derived_state[] = { 1, 2, 0, 0, 3, 0 }; +static const char *acgt[] = { "A", "C", "G", "T" }; +static double rho[] = { 0.01, 0.01, 0.01, 0.01, 0.01 }; +static double mu[] = { 0.01, 0.01, 0.01, 0.01, 0.01 }; + +static allele_t +acgt_code(char state) +{ + return state == 'A' ? 0 : state == 'C' ? 1 : state == 'G' ? 2 : 3; +} static void -test_matcher_indexes_errors(void) +test_allele_encoding(void) { int ret; - matcher_indexes_t mi; - tsk_table_collection_t tables; - allele_t bad_ancestral_state[] = { 0, 4, 2, 0 }; - allele_t bad_derived_state[] = { 1, 2, -1, 0, 3, 0 }; - - build_forest(&tables); - - ret = matcher_indexes_alloc(&mi, &tables, num_alleles, bad_ancestral_state, NULL, 0); - CU_ASSERT_EQUAL(ret, TSLS_ERR_BAD_ANCESTRAL_STATE); - matcher_indexes_free(&mi); - - ret = matcher_indexes_alloc(&mi, &tables, num_alleles, NULL, bad_derived_state, 0); - CU_ASSERT_EQUAL(ret, TSLS_ERR_BAD_MUTATION_DERIVED_STATE); - matcher_indexes_free(&mi); - - /* Default of 2 alleles */ - ret = matcher_indexes_alloc(&mi, &tables, NULL, ancestral_state, NULL, 0); - CU_ASSERT_EQUAL(ret, TSLS_ERR_BAD_ANCESTRAL_STATE); - matcher_indexes_free(&mi); - - tables.mutations.node[0] = 100; - ret = matcher_indexes_alloc(&mi, &tables, num_alleles, NULL, NULL, 0); - CU_ASSERT_EQUAL(ret, TSLS_ERR_BAD_MUTATION_NODE); - matcher_indexes_free(&mi); - tables.mutations.node[0] = 0; - - tables.mutations.site[0] = 100; - ret = matcher_indexes_alloc(&mi, &tables, num_alleles, NULL, NULL, 0); - CU_ASSERT_EQUAL(ret, TSLS_ERR_BAD_MUTATION_SITE); - matcher_indexes_free(&mi); - tables.mutations.site[0] = 0; - - tables.edges.child[0] = -1; - ret = matcher_indexes_alloc(&mi, &tables, num_alleles, NULL, NULL, 0); - CU_ASSERT_EQUAL(ret, TSLS_ERR_BAD_EDGE_NODE); - matcher_indexes_free(&mi); - tables.edges.child[0] = 0; - - tables.sites.position[0] = 1.5; - ret = matcher_indexes_alloc(&mi, &tables, num_alleles, NULL, NULL, 0); - CU_ASSERT_EQUAL(ret, TSLS_ERR_BAD_COORDINATE); - matcher_indexes_free(&mi); + tsk_treeseq_t ts; + tsls_hmm_t hmm; + tsk_size_t j; + tsk_mutation_t *mutation; + const char *too_few[] = { "A", "C" }; + /* Sites are A -> C, G; C -> A; G -> A; A -> T, A; T. Every site has at + * least two alleles. */ + tsk_size_t site_num_alleles[] = { 3, 2, 2, 2, 2 }; + + build_forest_ts(&ts, true); + + ret = tsls_hmm_alloc(&hmm, &ts, acgt, 4, rho, mu, DBL_MIN, 0); + CU_ASSERT_EQUAL_FATAL(ret, 0); + for (j = 0; j < tsk_treeseq_get_num_sites(&ts); j++) { + CU_ASSERT_EQUAL(hmm.num_alleles[j], 4); + CU_ASSERT_EQUAL(hmm.ancestral_state[j], + acgt_code(ts.tables->sites + .ancestral_state[ts.tables->sites.ancestral_state_offset[j]])); + } + for (j = 0; j < tsk_treeseq_get_num_mutations(&ts); j++) { + mutation = &ts.site_mutations_mem[j]; + CU_ASSERT_EQUAL( + hmm.derived_state[mutation->id], acgt_code(mutation->derived_state[0])); + } + tsls_hmm_free(&hmm); - tsk_table_collection_free(&tables); + /* Per-site encoding, as in tskit */ + ret = tsls_hmm_alloc(&hmm, &ts, NULL, 0, rho, mu, DBL_MIN, 0); + CU_ASSERT_EQUAL_FATAL(ret, 0); + for (j = 0; j < tsk_treeseq_get_num_sites(&ts); j++) { + CU_ASSERT_EQUAL(hmm.num_alleles[j], site_num_alleles[j]); + CU_ASSERT_EQUAL(hmm.ancestral_state[j], 0); + } + /* At the last site, the mutation to A is a back mutation */ + for (j = 0; j < ts.site_mutations_length[3]; j++) { + mutation = &ts.site_mutations[3][j]; + CU_ASSERT_EQUAL( + hmm.derived_state[mutation->id], mutation->derived_state[0] == 'A' ? 0 : 1); + } + tsls_hmm_free(&hmm); + + ret = tsls_hmm_alloc(&hmm, &ts, too_few, 2, rho, mu, DBL_MIN, 0); + CU_ASSERT_EQUAL(ret, TSLS_ERR_UNKNOWN_ALLELE); + tsls_hmm_free(&hmm); + + ret = tsls_hmm_alloc(&hmm, &ts, acgt, TSLS_MAX_ALLELES + 1, rho, mu, DBL_MIN, 0); + CU_ASSERT_EQUAL(ret, TSLS_ERR_TOO_MANY_ALLELES); + tsls_hmm_free(&hmm); + + tsk_treeseq_free(&ts); } static void test_forest_matching(void) { int ret; - tsk_table_collection_t tables; - matcher_indexes_t mi; - ancestor_matcher_t am; + tsk_treeseq_t ts; + tsls_hmm_t hmm; compressed_matrix_t fm, bm; - double rho[] = { 0.01, 0.01, 0.01, 0.01 }; - double mu[] = { 0.01, 0.01, 0.01, 0.01 }; /* The haplotype of node 1: C (under 3), C (ancestral), G, A (back mutation) */ allele_t h[] = { 1, 1, 2, 0 }; allele_t match[4]; - tsk_id_t left[4], right[4], parent[4]; + double left[4], right[4]; + tsk_id_t parent[4]; size_t path_length, j, k, l; double total; - build_forest(&tables); - ret = matcher_indexes_alloc(&mi, &tables, num_alleles, ancestral_state, derived_state, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - matcher_indexes_print_state(&mi, _devnull); - ret = ancestor_matcher_alloc(&am, &mi, rho, mu, DBL_MIN, TSLS_EXTENDED_CHECKS); + build_forest_ts(&ts, false); + ret = tsls_hmm_alloc(&hmm, &ts, acgt, 4, rho, mu, DBL_MIN, TSLS_EXTENDED_CHECKS); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = ancestor_matcher_find_path( - &am, 0, 4, h, match, &path_length, left, right, parent); + ret = tsls_hmm_find_path(&hmm, 0, 4, h, match, &path_length, left, right, parent); CU_ASSERT_EQUAL_FATAL(ret, 0); CU_ASSERT_EQUAL(path_length, 1); CU_ASSERT_EQUAL(parent[0], 1); @@ -144,19 +156,26 @@ test_forest_matching(void) for (j = 0; j < 4; j++) { CU_ASSERT_EQUAL(match[j], h[j]); } - CU_ASSERT(ancestor_matcher_get_mean_traceback_size(&am) > 0); - CU_ASSERT(ancestor_matcher_get_total_memory(&am) > 0); - ancestor_matcher_print_state(&am, _devnull); + CU_ASSERT(tsls_hmm_get_mean_traceback_size(&hmm) > 0); + CU_ASSERT(tsls_hmm_get_total_memory(&hmm) > 0); + tsls_hmm_print_state(&hmm, _devnull); + + /* A window within the right tree, which is found by seeking backwards */ + ret = tsls_hmm_find_path(&hmm, 2, 4, h, match, &path_length, left, right, parent); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL(path_length, 1); + CU_ASSERT_EQUAL(parent[0], 1); + CU_ASSERT_EQUAL(left[0], 6); + CU_ASSERT_EQUAL(right[0], 10); - ret = ancestor_matcher_find_path( - &am, 2, 1, h, match, &path_length, left, right, parent); + ret = tsls_hmm_find_path(&hmm, 2, 1, h, match, &path_length, left, right, parent); CU_ASSERT_EQUAL(ret, TSLS_ERR_BAD_SITE_RANGE); /* Viterbi with the values recorded */ ret = compressed_matrix_alloc(&fm, 4, 1024); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = ancestor_matcher_find_path_matrix( - &am, 0, 4, h, match, &path_length, left, right, parent, &fm); + ret = tsls_hmm_find_path_matrix( + &hmm, 0, 4, h, match, &path_length, left, right, parent, &fm); CU_ASSERT_EQUAL_FATAL(ret, 0); CU_ASSERT_EQUAL(path_length, 1); CU_ASSERT_EQUAL(parent[0], 1); @@ -171,9 +190,9 @@ test_forest_matching(void) CU_ASSERT_EQUAL_FATAL(ret, 0); ret = compressed_matrix_alloc(&bm, 4, 1024); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = ancestor_matcher_forward(&am, h, &fm); + ret = tsls_hmm_forward(&hmm, h, &fm); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = ancestor_matcher_backward(&am, h, fm.normalisation_factor, &bm); + ret = tsls_hmm_backward(&hmm, h, fm.normalisation_factor, &bm); CU_ASSERT_EQUAL_FATAL(ret, 0); CU_ASSERT(compressed_matrix_get_total_values(&fm) > 0); @@ -194,59 +213,54 @@ test_forest_matching(void) compressed_matrix_free(&fm); compressed_matrix_free(&bm); - ancestor_matcher_free(&am); - matcher_indexes_free(&mi); - tsk_table_collection_free(&tables); + tsls_hmm_free(&hmm); + tsk_treeseq_free(&ts); } static void test_bad_inputs(void) { int ret; - tsk_table_collection_t tables; - matcher_indexes_t mi; - ancestor_matcher_t am; + tsk_treeseq_t ts; + tsls_hmm_t hmm; compressed_matrix_t cm; - double rho[] = { 0.01, 0.01, 0.01, 0.01 }; - double mu[] = { 0, 0, 0, 0 }; + double zero_mu[] = { 0, 0, 0, 0 }; double norm[] = { 1, 1, 0, 1 }; allele_t h[] = { 1, 1, 2, 0 }; allele_t bad_h[] = { 1, 4, 2, 0 }; allele_t impossible_h[] = { 3, 3, 3, 3 }; allele_t match[4]; - tsk_id_t left[4], right[4], parent[4]; + double left[4], right[4]; + tsk_id_t parent[4]; size_t path_length; - build_forest(&tables); - ret = matcher_indexes_alloc(&mi, &tables, num_alleles, ancestral_state, derived_state, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = ancestor_matcher_alloc(&am, &mi, rho, mu, DBL_MIN, 0); + build_forest_ts(&ts, false); + ret = tsls_hmm_alloc(&hmm, &ts, acgt, 4, rho, zero_mu, DBL_MIN, 0); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = ancestor_matcher_find_path( - &am, 0, 4, bad_h, match, &path_length, left, right, parent); + ret = tsls_hmm_find_path( + &hmm, 0, 4, bad_h, match, &path_length, left, right, parent); CU_ASSERT_EQUAL(ret, TSLS_ERR_BAD_HAPLOTYPE_ALLELE); - ret = ancestor_matcher_find_path( - &am, 0, 4, impossible_h, match, &path_length, left, right, parent); + ret = tsls_hmm_find_path( + &hmm, 0, 4, impossible_h, match, &path_length, left, right, parent); CU_ASSERT_EQUAL(ret, TSLS_ERR_MATCH_IMPOSSIBLE_EXTREME_MUTATION_PROBA); ret = compressed_matrix_alloc(&cm, 4, 1024); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = ancestor_matcher_forward(&am, impossible_h, &cm); + ret = tsls_hmm_forward(&hmm, impossible_h, &cm); CU_ASSERT_EQUAL(ret, TSLS_ERR_MATCH_IMPOSSIBLE); - ret = ancestor_matcher_backward(&am, h, norm, &cm); + ret = tsls_hmm_backward(&hmm, h, norm, &cm); CU_ASSERT_EQUAL(ret, TSLS_ERR_BAD_NORMALISATION_FACTOR); compressed_matrix_free(&cm); ret = compressed_matrix_alloc(&cm, 3, 1024); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = ancestor_matcher_forward(&am, h, &cm); + ret = tsls_hmm_forward(&hmm, h, &cm); CU_ASSERT_EQUAL(ret, TSLS_ERR_BAD_SITE_RANGE); compressed_matrix_free(&cm); - ancestor_matcher_free(&am); - matcher_indexes_free(&mi); - tsk_table_collection_free(&tables); + tsls_hmm_free(&hmm); + tsk_treeseq_free(&ts); } static void @@ -296,7 +310,7 @@ main(int argc, char **argv) CU_pTest test; CU_pSuite suite; CU_TestInfo tests[] = { - { "test_matcher_indexes_errors", test_matcher_indexes_errors }, + { "test_allele_encoding", test_allele_encoding }, { "test_forest_matching", test_forest_matching }, { "test_bad_inputs", test_bad_inputs }, { "test_strerror", test_strerror }, diff --git a/tsls/lib/tsls.h b/tsls/lib/tsls.h index 506a238..739b423 100644 --- a/tsls/lib/tsls.h +++ b/tsls/lib/tsls.h @@ -47,9 +47,11 @@ #define TSLS_PRUNE_TOP_N 2 typedef int8_t allele_t; -typedef tsk_id_t coordinate_t; -/* Allow the CPython accessor AncestorMatcher.total_memory to read the +/* The maximum number of distinct alleles at a site */ +#define TSLS_MAX_ALLELES INT8_MAX + +/* Allow the CPython accessor Hmm.total_memory to read the * current allocator usage accurately while other threads may be allocating. * Only total_size is atomic so it can be observed safely from the getter * without data races. */ @@ -69,44 +71,12 @@ extern int tsi_blkalloc_init(tsi_blkalloc_t *self, size_t chunk_size); extern void *tsi_blkalloc_get(tsi_blkalloc_t *self, size_t size); extern void tsi_blkalloc_free(tsi_blkalloc_t *self); -typedef struct { - coordinate_t left; - coordinate_t right; - tsk_id_t parent; - tsk_id_t child; -} edge_t; - -typedef struct _mutation_list_node_t { - tsk_id_t node; - allele_t derived_state; - struct _mutation_list_node_t *next; -} mutation_list_node_t; - typedef struct { int32_t size; tsk_id_t *node; int8_t *recombination_required; } node_state_list_t; -/* The read-only information about a tree sequence needed for matching, - * which can be shared between matchers in different threads. */ -typedef struct { - tsk_flags_t flags; - size_t num_sites; - size_t num_nodes; - size_t num_mutations; - size_t num_edges; - struct { - coordinate_t *position; - mutation_list_node_t **mutations; - tsk_size_t *num_alleles; - allele_t *ancestral_state; - } sites; - edge_t *left_index_edges; - edge_t *right_index_edges; - tsi_blkalloc_t allocator; -} matcher_indexes_t; - /* The compressed output of the forward and backward algorithms. At each * site we store a list of (node, value) pairs: the value for a node in the * tree is the value of its nearest ancestor (including itself) in this list. @@ -122,11 +92,20 @@ typedef struct { tsi_blkalloc_t allocator; } compressed_matrix_t; +/* The Li and Stephens HMM on a tree sequence. The tree sequence is borrowed, + * and must not be changed or freed while the HMM is in use. */ typedef struct { int flags; - const matcher_indexes_t *matcher_indexes; + const tsk_treeseq_t *tree_sequence; size_t num_nodes; size_t num_sites; + const double *sites_position; + double sequence_length; + /* The integer encoding of the alleles: the number of alleles and the + * ancestral state of each site, and the derived state of each mutation */ + tsk_size_t *num_alleles; + allele_t *ancestral_state; + allele_t *derived_state; /* Input LS model rates */ double likelihood_threshold; int pruning_mode; @@ -168,39 +147,33 @@ typedef struct { double *values_tmp; size_t traceback_block_size; size_t traceback_realloc_size; -} ancestor_matcher_t; - -int matcher_indexes_alloc(matcher_indexes_t *self, const tsk_table_collection_t *tables, - const tsk_size_t *num_alleles, const allele_t *ancestral_state, - const allele_t *derived_state, tsk_flags_t options); -int matcher_indexes_print_state(const matcher_indexes_t *self, FILE *out); -int matcher_indexes_free(matcher_indexes_t *self); +} tsls_hmm_t; int compressed_matrix_alloc( compressed_matrix_t *self, size_t num_sites, size_t block_size); int compressed_matrix_free(compressed_matrix_t *self); size_t compressed_matrix_get_total_values(const compressed_matrix_t *self); -int ancestor_matcher_alloc(ancestor_matcher_t *self, - const matcher_indexes_t *matcher_indexes, const double *recombination_rate, +int tsls_hmm_alloc(tsls_hmm_t *self, const tsk_treeseq_t *tree_sequence, + const char **alleles, tsk_size_t num_alleles, const double *recombination_rate, const double *mismatch_rate, double likelihood_threshold, int flags); -int ancestor_matcher_set_pruning(ancestor_matcher_t *self, int mode, double parameter); -int ancestor_matcher_free(ancestor_matcher_t *self); -int ancestor_matcher_find_path(ancestor_matcher_t *self, tsk_id_t start, tsk_id_t end, +int tsls_hmm_set_pruning(tsls_hmm_t *self, int mode, double parameter); +int tsls_hmm_free(tsls_hmm_t *self); +int tsls_hmm_find_path(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, + const allele_t *haplotype, allele_t *matched_haplotype, size_t *path_length, + double *path_left, double *path_right, tsk_id_t *path_parent); +int tsls_hmm_find_path_matrix(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, const allele_t *haplotype, allele_t *matched_haplotype, size_t *path_length, - tsk_id_t *path_left, tsk_id_t *path_right, tsk_id_t *path_parent); -int ancestor_matcher_find_path_matrix(ancestor_matcher_t *self, tsk_id_t start, - tsk_id_t end, const allele_t *haplotype, allele_t *matched_haplotype, - size_t *path_length, tsk_id_t *path_left, tsk_id_t *path_right, - tsk_id_t *path_parent, compressed_matrix_t *output); -int ancestor_matcher_forward(ancestor_matcher_t *self, const allele_t *haplotype, + double *path_left, double *path_right, tsk_id_t *path_parent, compressed_matrix_t *output); -int ancestor_matcher_backward(ancestor_matcher_t *self, const allele_t *haplotype, +int tsls_hmm_forward( + tsls_hmm_t *self, const allele_t *haplotype, compressed_matrix_t *output); +int tsls_hmm_backward(tsls_hmm_t *self, const allele_t *haplotype, const double *forward_normalisation_factor, compressed_matrix_t *output); -int ancestor_matcher_print_state(ancestor_matcher_t *self, FILE *out); -double ancestor_matcher_get_mean_traceback_size(ancestor_matcher_t *self); -double ancestor_matcher_get_mean_num_distinct_values(ancestor_matcher_t *self); -size_t ancestor_matcher_get_total_memory(ancestor_matcher_t *self); +int tsls_hmm_print_state(tsls_hmm_t *self, FILE *out); +double tsls_hmm_get_mean_traceback_size(tsls_hmm_t *self); +double tsls_hmm_get_mean_num_distinct_values(tsls_hmm_t *self); +size_t tsls_hmm_get_total_memory(tsls_hmm_t *self); #define tsi_safe_free(pointer) \ do { \ diff --git a/tsls/setup.py b/tsls/setup.py index 29ae621..c673f03 100644 --- a/tsls/setup.py +++ b/tsls/setup.py @@ -13,13 +13,20 @@ includes = ["lwt_interface", libdir, tskroot, tskdir, kasdir] tsls_source_files = [ - "matcher.c", + "hmm.c", "err.c", "tsi_blkalloc.c", ] -# The matcher only needs the tskit core and the table code used by the -# LightweightTableCollection interface. -tsk_source_files = ["core.c", "tables.c"] +# The HMM works on a tsk_treeseq_t, so we need all of tskit. +tsk_source_files = [ + "core.c", + "tables.c", + "trees.c", + "genotypes.c", + "stats.c", + "convert.c", + "haplotype_matching.c", +] kas_source_files = ["kastore.c"] sources = ( diff --git a/tsls/tests/test_matching.py b/tsls/tests/test_matching.py index df4282d..b951491 100644 --- a/tsls/tests/test_matching.py +++ b/tsls/tests/test_matching.py @@ -255,11 +255,11 @@ def queries(ts, A, present): def make_matcher(ts, rho, mu, weight_by_n, likelihood_threshold=None): m = ts.num_sites - mi = tsls.MatcherIndexes(ts, alleles=ALLELES) - am = tsls.AncestorMatcher( - mi, + am = tsls.Hmm( + ts, np.full(m, rho), np.full(m, mu), + alleles=ALLELES, weight_by_n=weight_by_n, likelihood_threshold=likelihood_threshold, extended_checks=True, @@ -268,16 +268,20 @@ def make_matcher(ts, rho, mu, weight_by_n, likelihood_threshold=None): return am, n -def path_by_site(ts, left, right, parent): +def path_by_site(ts, left, right, parent, start=0, end=None): """ - Return the node copied at each site, from a path in any order. + Return the node copied at each site in [start, end), from a path in any + order. """ + if end is None: + end = ts.num_sites index = np.argsort(left) left, right, parent = left[index], right[index], parent[index] - assert left[0] == 0 - assert right[-1] == ts.sequence_length + position = ts.sites_position + assert left[0] == (0 if start == 0 else position[start]) + assert right[-1] == (ts.sequence_length if end == ts.num_sites else position[end]) np.testing.assert_array_equal(left[1:], right[:-1]) - segment = np.searchsorted(left, ts.sites_position, side="right") - 1 + segment = np.searchsorted(left, position[start:end], side="right") - 1 return parent[segment] @@ -460,22 +464,101 @@ def test_start_end(self): assert list(right) == [ts.sites_position[-1]] np.testing.assert_array_equal(match[1:-1], h[1:-1]) + @pytest.mark.parametrize("name", list(EXAMPLES)) + def test_windows(self, name): + # The first tree of a window is found by seeking from whichever end + # of the sequence is closer, so check windows starting and ending + # throughout the sequence against the dense optimum for the window. + ts = EXAMPLES[name]() + A, present, continuous = dense.node_states(ts, ALLELES) + rho, mu = 0.05, 0.02 + am, n = make_matcher(ts, rho, mu, True) + m = ts.num_sites + for h in queries(ts, A, present)[:4]: + for start in range(m): + for end in range(start + 1, m + 1): + match = np.zeros(m, dtype=np.int8) + left, right, parent = am.find_path(h, start, end, match) + path = path_by_site(ts, left, right, parent, start, end) + sites = np.arange(start, end) + np.testing.assert_array_equal(match[start:end], A[sites, path]) + window = (A[start:end], present[start:end]) + cont = continuous[start:end].copy() + cont[0] = present[start] + rho_a = np.full(end - start, rho) + mu_a = np.full(end - start, mu) + args = (*window, cont, h[start:end], rho_a, mu_a, len(ALLELES), n) + value = dense.path_log_likelihood(*args, path) + optimum = dense.viterbi_log_likelihood(*args) + assert value == pytest.approx(optimum, rel=1e-9) + class TestInterface: - def test_bad_alleles(self): - with pytest.raises(ValueError, match="Unknown"): - tsls.MatcherIndexes(single_tree_ts(), alleles=["A", "C"]) + def test_unknown_allele(self): + ts = single_tree_ts() + m = ts.num_sites + with pytest.raises(tsls.LibraryError, match="not in the list"): + tsls.Hmm(ts, np.zeros(m), np.zeros(m), alleles=["A", "C"]) + + def test_duplicate_alleles(self): + ts = single_tree_ts() + m = ts.num_sites + with pytest.raises(ValueError, match="unique"): + tsls.Hmm(ts, np.zeros(m), np.zeros(m), alleles=["A", "A"]) def test_default_binary_alleles(self): tables = tskit.Tree.generate_balanced(2, span=4).tree_sequence.dump_tables() add_mutations(tables, [(1, "0", [(0, "1")])]) ts = finalise(tables) - mi = tsls.MatcherIndexes(ts) - am = tsls.AncestorMatcher(mi, np.full(1, 1e-3), np.full(1, 1e-3)) + am = tsls.Hmm(ts, np.full(1, 1e-3), np.full(1, 1e-3)) match = np.zeros(1, dtype=np.int8) _, _, parent = am.find_path(np.array([1], dtype=np.int8), 0, 1, match) assert list(parent) == [0] + @pytest.mark.parametrize("name", list(EXAMPLES)) + def test_per_site_alleles(self, name): + # By default the alleles are numbered at each site as tskit does, so + # the matched haplotypes are tskit's genotypes. + ts = EXAMPLES[name]() + m = ts.num_sites + am = tsls.Hmm(ts, np.full(m, 1e-4), np.full(m, 1e-3), extended_checks=True) + G = ts.genotype_matrix(isolated_as_missing=False) + for j in range(ts.num_samples): + h = G[:, j].astype(np.int8) + match = np.zeros(m, dtype=np.int8) + am.find_path(h, 0, m, match) + np.testing.assert_array_equal(match, h) + + def test_site_without_mutations(self): + # As in tsinfer's ancestor matching, a haplotype can carry an allele + # at a site that has no mutations in the tree sequence. + tables = tskit.Tree.generate_balanced(2, span=4).tree_sequence.dump_tables() + add_mutations(tables, [(1, "0", [(0, "1")]), (2, "0", [])]) + ts = finalise(tables) + am = tsls.Hmm(ts, np.full(2, 1e-3), np.full(2, 1e-3)) + match = np.zeros(2, dtype=np.int8) + _, _, parent = am.find_path(np.array([1, 1], dtype=np.int8), 0, 2, match) + assert list(parent) == [0] + assert list(match) == [1, 0] + with pytest.raises(tsls.LibraryError, match="Haplotype allele"): + am.find_path(np.array([1, 2], dtype=np.int8), 0, 2, match) + + def test_shared_tree_sequence(self): + ts = two_tree_ts() + tts = tsls.TreeSequence(ts) + assert tts.num_nodes == ts.num_nodes + assert tts.num_sites == ts.num_sites + m = ts.num_sites + am1 = tsls.Hmm(tts, np.full(m, 0.01), np.full(m, 0.01), alleles=ALLELES) + am2 = tsls.Hmm(ts, np.full(m, 0.01), np.full(m, 0.01), alleles=ALLELES) + assert am1.tree_sequence is tts + h = np.zeros(m, dtype=np.int8) + for x, y in zip( + am1.find_path(h, 0, m, np.zeros(m, dtype=np.int8)), + am2.find_path(h, 0, m, np.zeros(m, dtype=np.int8)), + ): + np.testing.assert_array_equal(x, y) + @pytest.mark.parametrize("state", [5, -2]) def test_bad_haplotype_allele(self, state): ts = single_tree_ts() @@ -498,12 +581,11 @@ def test_bad_site_range(self, start, end): am.find_path(h, start, end, match) def test_bad_lengths(self): - ts = single_tree_ts() - mi = tsls.MatcherIndexes(ts, alleles=ALLELES) + ts = tsls.TreeSequence(single_tree_ts()) with pytest.raises(ValueError, match="recombination"): - tsls.AncestorMatcher(mi, np.zeros(1), np.zeros(ts.num_sites)) + tsls.Hmm(ts, np.zeros(1), np.zeros(ts.num_sites)) with pytest.raises(ValueError, match="mismatch"): - tsls.AncestorMatcher(mi, np.zeros(ts.num_sites), np.zeros(1)) + tsls.Hmm(ts, np.zeros(ts.num_sites), np.zeros(1)) am, _ = make_matcher(ts, 1e-2, 1e-2, True) with pytest.raises(ValueError): am.forward_matrix(np.zeros(1, dtype=np.int8)) @@ -529,10 +611,17 @@ def test_match_impossible(self): am.forward_matrix(h) def test_non_integer_coordinates(self): - tables = tskit.Tree.generate_balanced(2, span=4).tree_sequence.dump_tables() - tables.sites.add_row(position=1.5, ancestral_state="0") - with pytest.raises(tsls.LibraryError, match="integers"): - tsls.MatcherIndexes(finalise(tables)) + # Node 0 on [0, 1.5), and node 1 on [1.5, 3.25) + tables = tskit.Tree.generate_balanced(2, span=3.25).tree_sequence.dump_tables() + add_mutations(tables, [(0.5, "0", [(0, "1")]), (2.75, "0", [(1, "1")])]) + ts = finalise(tables) + am = tsls.Hmm(ts, np.full(2, 0.4), np.full(2, 1e-3)) + match = np.zeros(2, dtype=np.int8) + left, right, parent = am.find_path(np.array([1, 1], dtype=np.int8), 0, 2, match) + assert left.dtype == np.float64 + assert list(parent) == [1, 0] + assert list(left) == [2.75, 0] + assert list(right) == [3.25, 2.75] def test_memory_stats(self): ts = two_tree_ts() @@ -577,9 +666,13 @@ def test_compressed_matrix(self): class TestPruning: def run_all(self, ts, h, **kwargs): m = ts.num_sites - mi = tsls.MatcherIndexes(ts, alleles=ALLELES) - am = tsls.AncestorMatcher( - mi, np.full(m, 0.05), np.full(m, 0.02), extended_checks=True, **kwargs + am = tsls.Hmm( + ts, + np.full(m, 0.05), + np.full(m, 0.02), + alleles=ALLELES, + extended_checks=True, + **kwargs, ) match = np.zeros(m, dtype=np.int8) path = am.find_path(h, 0, m, match) @@ -653,17 +746,15 @@ def test_count_distinct(self): ) def test_bad_pruning(self, pruning): ts = single_tree_ts() - mi = tsls.MatcherIndexes(ts, alleles=ALLELES) m = ts.num_sites with pytest.raises(tsls.LibraryError, match="pruning"): - tsls.AncestorMatcher(mi, np.zeros(m), np.zeros(m), pruning=pruning) + tsls.Hmm(ts, np.zeros(m), np.zeros(m), alleles=ALLELES, pruning=pruning) def test_unknown_pruning_mode(self): ts = single_tree_ts() - mi = tsls.MatcherIndexes(ts, alleles=ALLELES) m = ts.num_sites with pytest.raises(ValueError, match="Unknown pruning"): - tsls.AncestorMatcher(mi, np.zeros(m), np.zeros(m), pruning=("x", 1)) + tsls.Hmm(ts, np.zeros(m), np.zeros(m), pruning=("x", 1)) class TestViterbiMatrix: diff --git a/tsls/tsls/__init__.py b/tsls/tsls/__init__.py index e5e0f5d..7debce0 100644 --- a/tsls/tsls/__init__.py +++ b/tsls/tsls/__init__.py @@ -5,9 +5,8 @@ from _tsls import LibraryError, MatchImpossible # noqa: F401 from .matching import ( # noqa: F401 - AncestorMatcher, CompressedMatrix, - MatcherIndexes, - encode_alleles, + Hmm, + TreeSequence, load_compressed_matrix, ) diff --git a/tsls/tsls/matching.py b/tsls/tsls/matching.py index 4a241a2..7d2adab 100644 --- a/tsls/tsls/matching.py +++ b/tsls/tsls/matching.py @@ -7,66 +7,19 @@ import _tsls -def encode_alleles(ts, alleles): +class TreeSequence(_tsls.TreeSequence): """ - Return the (ancestral_state, derived_state) arrays for the specified tree - sequence, encoding each state as its index in the specified list of - alleles. Raise a ValueError if any state is not in the list. + A tskit tree sequence in the form used by the HMM. It can be shared + between HMMs, including HMMs in different threads, so that it is only + built once. """ - alleles = list(alleles) - if len(alleles) > np.iinfo(np.int8).max: - raise ValueError("Too many alleles") - allele_index = {allele: j for j, allele in enumerate(alleles)} - if len(allele_index) != len(alleles): - raise ValueError("Alleles must be unique") - def encode(states, name): - unique_states, inverse = np.unique(np.asarray(states), return_inverse=True) - code = np.zeros(len(unique_states), dtype=np.int8) - for j, state in enumerate(unique_states): - if state not in allele_index: - raise ValueError(f"Unknown {name}: {state!r}") - code[j] = allele_index[state] - return code[inverse] - - ancestral_state = encode(ts.sites_ancestral_state, "ancestral state") - derived_state = encode(ts.mutations_derived_state, "derived state") - return ancestral_state, derived_state - - -class MatcherIndexes(_tsls.MatcherIndexes): - """ - The information about a tree sequence needed for matching, which can be - shared between matchers (including in different threads). - - The states of the HMM at a given position are the nodes in the tree at - that position: nodes with a parent or at least one child. Isolated - nodes cannot be copied. Coordinates must be integers. - - By default, the ancestral state of every site is 0 and every mutation has - derived state 1. If ``alleles`` is specified, states are instead encoded - as their index in this list, and ``num_alleles`` defaults to its length. - """ - - def __init__(self, ts, *, alleles=None, num_alleles=None): - ancestral_state = None - derived_state = None - if alleles is not None: - ancestral_state, derived_state = encode_alleles(ts, alleles) - if num_alleles is None: - num_alleles = np.full(ts.num_sites, len(alleles), dtype=np.uint64) + def __init__(self, ts): tables = ts.dump_tables() ll_tables = _tsls.LightweightTableCollection(tables.sequence_length) ll_tables.fromdict(tables.asdict()) self.ts = ts - self.sequence_length = ts.sequence_length - self.num_sites = ts.num_sites - super().__init__( - ll_tables, - num_alleles=num_alleles, - ancestral_state=ancestral_state, - derived_state=derived_state, - ) + super().__init__(ll_tables) class CompressedMatrix: @@ -181,11 +134,23 @@ def load_compressed_matrix(path, ts): ) -class AncestorMatcher(_tsls.AncestorMatcher): +class Hmm(_tsls.Hmm): """ Runs the Viterbi, forward and backward algorithms for haplotypes against - the tree sequence in the specified MatcherIndexes, using the specified - per-site recombination and mismatch probabilities. + the specified tree sequence (a tskit.TreeSequence, or a tsls.TreeSequence + to share one between HMMs), using the specified per-site recombination + and mismatch probabilities. + + The states of the HMM at a given position are the nodes in the tree at + that position: nodes with a parent or at least one child. Isolated + nodes cannot be copied. + + Haplotypes are integer encoded. By default the alleles are numbered at + each site as tskit does: the ancestral state is 0 and derived states are + numbered from 1 in order of their first appearance at the site. Every + site has at least two alleles, so that 1 can be matched at a site with + no mutations (as in tsinfer, for an ancestor's focal sites). If + ``alleles`` is specified, every state is instead its index in this list. After normalisation at each site, values below likelihood_threshold are set to it. The experimental ``pruning`` option also sets the values of @@ -200,17 +165,25 @@ class AncestorMatcher(_tsls.AncestorMatcher): def __init__( self, - matcher_indexes, + ts, recombination, mismatch, *, + alleles=None, likelihood_threshold=None, weight_by_n=None, extended_checks=False, pruning=None, count_distinct=False, ): + if not isinstance(ts, TreeSequence): + ts = TreeSequence(ts) kwargs = {"extended_checks": extended_checks, "count_distinct": count_distinct} + if alleles is not None: + alleles = list(alleles) + if len(set(alleles)) != len(alleles): + raise ValueError("Alleles must be unique") + kwargs["alleles"] = alleles if pruning is not None: mode, parameter = pruning if mode not in self.PRUNING_MODES: @@ -221,13 +194,9 @@ def __init__( kwargs["likelihood_threshold"] = likelihood_threshold if weight_by_n is not None: kwargs["weight_by_n"] = weight_by_n - self.ts = matcher_indexes.ts - super().__init__( - matcher_indexes, - recombination=recombination, - mismatch=mismatch, - **kwargs, - ) + self.tree_sequence = ts + self.ts = ts.ts + super().__init__(ts, recombination=recombination, mismatch=mismatch, **kwargs) def find_path(self, h, start, end, match_out): """ diff --git a/verification/verification.py b/verification/verification.py index 257e06b..66c907b 100644 --- a/verification/verification.py +++ b/verification/verification.py @@ -146,15 +146,16 @@ def make_matcher( likelihood_threshold=None, pruning=None, count_distinct=False, - matcher_indexes=None, + tree_sequence=None, ): m = ts.num_sites - if matcher_indexes is None: - matcher_indexes = tsls.MatcherIndexes(ts, alleles=ALLELES) - return tsls.AncestorMatcher( - matcher_indexes, + if tree_sequence is None: + tree_sequence = tsls.TreeSequence(ts) + return tsls.Hmm( + tree_sequence, np.full(m, rho), np.full(m, mu), + alleles=ALLELES, weight_by_n=weight_by_n, likelihood_threshold=likelihood_threshold, pruning=pruning, From 912d4a44e46ba59fccc23a301aaa44bea7f503ea Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:08:58 +0100 Subject: [PATCH 05/25] Reduce the callgrind profiling set to a few minutes The first harness took over 20 minutes per engine, mostly in setup: under callgrind, loading a large reference tree sequence and building the engine took 4-5 minutes per workload before any matching. Now: - Setup runs uninstrumented (--instr-atstart=no), and the driver switches instrumentation on with callgrind_control just before the queries. Setup for g1485 fell from 235 s to 35 s. Collection is still limited to find_path, so the counts are the same. - Fewer queries, sized at about 2G tsinfer-hmm instructions per workload (matching runs at about 36M instructions/s under callgrind): 20 g84, 20 g1000, 10 g1483, 3 g1484 and 5 g1485 ancestors, and 1 sample restricted to the middle 100,000 of its sites (a whole sample is 10G). - tsinfer-hmm's counts never change, so they are measured once and reused from callgrind/tsinfer/; each run measures only tsls. - Paths are saved next to the callgrind files and compared between the engines. A run measuring both engines takes 2.6 minutes with 4 jobs; tsls alone is about half that. --- .gitignore | 1 + benchmarking/callgrind.py | 181 ++++++++++++++++++++++++-------------- 2 files changed, 118 insertions(+), 64 deletions(-) diff --git a/.gitignore b/.gitignore index 3a9297e..05ff0d8 100644 --- a/.gitignore +++ b/.gitignore @@ -224,6 +224,7 @@ __marimo__/ # Benchmarking data and results benchmarking/data/ benchmarking/results/ +benchmarking/callgrind/ # Verification output verification/tmp__NOBACKUP__/ diff --git a/benchmarking/callgrind.py b/benchmarking/callgrind.py index 9d17c1b..3e208e1 100644 --- a/benchmarking/callgrind.py +++ b/benchmarking/callgrind.py @@ -2,13 +2,20 @@ Count the instructions the HMM engines execute on the tsinfer workloads. ``run`` matches the first queries of a workload with one engine. ``measure`` -runs ``run`` under callgrind, with collection limited to the engine's -find_path call, so that the counts are deterministic and leave out setup -(loading the data and building the engine). It prints a markdown table of -total instructions (Ir) for each workload and engine. +runs ``run`` under callgrind for each workload, with collection limited to +the engine's find_path call, so that the counts are deterministic and leave +out setup (loading the data and building the engine). Setup also runs +uninstrumented, which makes it many times faster. It prints a markdown +table of total instructions (Ir) for each workload, comparing tsls with +tsinfer-hmm. + +The workloads are sized so that measuring tsls takes a few minutes. The +tsinfer-hmm counts do not change, so they are measured once and reused from +the baseline directory. """ import concurrent.futures +import os import pathlib import subprocess import sys @@ -23,17 +30,20 @@ # The C function that runs one Viterbi match in each engine. Collection is # switched on only inside these. FIND_PATH_FUNCTIONS = { - "tsinfer": ["ancestor_matcher_find_path"], - "tsls": ["tsls_hmm_find_path", "ancestor_matcher_find_path"], + "tsinfer": "ancestor_matcher_find_path", + "tsls": "tsls_hmm_find_path", } -DEFAULT_WORKLOADS = [ - ("ancestors_g84", 20), - ("ancestors_g1000", 20), - ("ancestors_g1483", 20), - ("ancestors_g1484", 20), - ("ancestors_g1485", 20), - ("samples", 1), +# (workload, number of queries, maximum number of sites per query). Each is +# about 2G instructions in tsinfer-hmm. A whole sample is about 10G, so the +# sample is restricted to the middle of its window. +WORKLOADS = [ + ("ancestors_g84", 20, None), + ("ancestors_g1000", 20, None), + ("ancestors_g1483", 10, None), + ("ancestors_g1484", 3, None), + ("ancestors_g1485", 5, None), + ("samples", 1, 100_000), ] @@ -47,6 +57,16 @@ def make_engine(name, ts, model): ) +def clip_window(start, end, max_sites): + """ + Return the middle max_sites sites of the window [start, end). + """ + if max_sites is None or end - start <= max_sites: + return start, end + start = start + (end - start - max_sites) // 2 + return start, start + max_sites + + @click.group() def cli(): pass @@ -56,18 +76,32 @@ def cli(): @click.argument("workload_dir", type=click.Path(exists=True, file_okay=False)) @click.argument("engine_name", type=click.Choice(list(FIND_PATH_FUNCTIONS))) @click.option("-n", "--num-queries", type=int, default=20) -def run(workload_dir, engine_name, num_queries): +@click.option("--max-sites", type=int, default=None) +@click.option( + "--instrument", + is_flag=True, + help="Switch on callgrind instrumentation after setup", +) +def run(workload_dir, engine_name, num_queries, max_sites, instrument): """ Match the first NUM_QUERIES queries in WORKLOAD_DIR with the specified - engine, printing the parents of each path. + engine, each over at most MAX_SITES sites in the middle of its window, + printing the parents of each path. """ workload_dir = pathlib.Path(workload_dir) model = queries.load_model(workload_dir) ts = run_module.load_ts(workload_dir / "reference.trees") H, _, df = queries.load_queries(workload_dir) engine = make_engine(engine_name, ts, model) + if instrument: + subprocess.run( + ["callgrind_control", "-i", "on", str(os.getpid())], + check=True, + stdout=subprocess.DEVNULL, + ) for query_id in df.query_id[:num_queries]: start, end = queries.query_window(H, query_id) + start, end = clip_window(start, end, max_sites) _, _, parent, _, _ = engine.match(H[query_id], start, end) print(query_id, " ".join(str(p) for p in np.asarray(parent))) @@ -75,32 +109,42 @@ def run(workload_dir, engine_name, num_queries): def read_total_ir(path): with open(path) as f: for line in f: - if line.startswith("totals:") or line.startswith("summary:"): + if line.startswith(("totals:", "summary:")): return int(line.split()[1]) raise ValueError(f"No totals in {path}") -def callgrind(workload_dir, engine_name, num_queries, out_dir): - workload = pathlib.Path(workload_dir).name - out_file = pathlib.Path(out_dir) / f"callgrind.{workload}.{engine_name}.out" - toggles = [f"--toggle-collect={f}" for f in FIND_PATH_FUNCTIONS[engine_name]] +def out_path(out_dir, workload, num_queries, max_sites, engine_name): + sites = "" if max_sites is None else f".s{max_sites}" + return pathlib.Path(out_dir) / ( + f"callgrind.{workload}.n{num_queries}{sites}.{engine_name}.out" + ) + + +def callgrind(data_dir, workload, num_queries, max_sites, engine_name, out_file): cmd = [ "valgrind", "--tool=callgrind", + "--instr-atstart=no", "--collect-atstart=no", - *toggles, + f"--toggle-collect={FIND_PATH_FUNCTIONS[engine_name]}", f"--callgrind-out-file={out_file}", sys.executable, __file__, "run", - str(workload_dir), + str(pathlib.Path(data_dir) / workload), engine_name, "-n", str(num_queries), + "--instrument", ] - result = subprocess.run(cmd, capture_output=True, text=True) + if max_sites is not None: + cmd += ["--max-sites", str(max_sites)] + result = subprocess.run(cmd, check=False, capture_output=True, text=True) if result.returncode != 0: raise RuntimeError(f"{' '.join(cmd)} failed:\n{result.stderr[-2000:]}") + paths_file = out_file.with_suffix(".paths") + paths_file.write_text(result.stdout) return read_total_ir(out_file), result.stdout @@ -112,53 +156,62 @@ def callgrind(workload_dir, engine_name, num_queries, out_dir): ) @click.option("-w", "--workload", "workload_names", multiple=True) @click.option( - "-e", - "--engine", - "engine_names", - multiple=True, - type=click.Choice(list(FIND_PATH_FUNCTIONS)), + "-o", + "--out-dir", + type=click.Path(file_okay=False), + default="callgrind", + help="Directory for the tsls callgrind output files", +) +@click.option( + "--baseline-dir", + type=click.Path(file_okay=False), + default="callgrind/tsinfer", + help="Directory for the tsinfer-hmm callgrind output files, which are reused", ) -@click.option("-o", "--out-dir", type=click.Path(file_okay=False), default="callgrind") @click.option("-j", "--jobs", type=int, default=3) -def measure(data_dir, workload_names, engine_names, out_dir, jobs): +def measure(data_dir, workload_names, out_dir, baseline_dir, jobs): """ - Count the instructions in find_path for each workload and engine. The - callgrind output files are kept in OUT_DIR for callgrind_annotate. + Count the instructions in find_path for each workload, for tsls and + (unless already measured) tsinfer-hmm. The callgrind output files are + kept for callgrind_annotate. """ - workloads = DEFAULT_WORKLOADS + workloads = WORKLOADS if len(workload_names) > 0: - workloads = [w for w in DEFAULT_WORKLOADS if w[0] in workload_names] - if len(engine_names) == 0: - engine_names = list(FIND_PATH_FUNCTIONS) - pathlib.Path(out_dir).mkdir(parents=True, exist_ok=True) - tasks = [ - (workload, num_queries, engine) - for workload, num_queries in workloads - for engine in engine_names - ] + workloads = [w for w in WORKLOADS if w[0] in workload_names] + tasks = [] + for workload, num_queries, max_sites in workloads: + for engine_name, directory in [("tsinfer", baseline_dir), ("tsls", out_dir)]: + pathlib.Path(directory).mkdir(parents=True, exist_ok=True) + path = out_path(directory, workload, num_queries, max_sites, engine_name) + if engine_name == "tsinfer" and path.exists(): + continue + tasks.append((workload, num_queries, max_sites, engine_name, path)) + # Run the largest workloads first, so that they don't finish last + tasks.sort(key=lambda task: task[0] not in ("samples", "ancestors_g1485")) with concurrent.futures.ThreadPoolExecutor(jobs) as executor: - futures = { - task: executor.submit( - callgrind, pathlib.Path(data_dir) / task[0], task[2], task[1], out_dir - ) - for task in tasks - } - results = {task: future.result() for task, future in futures.items()} - - header = "| workload | queries | " + " | ".join(engine_names) + " |" - if len(engine_names) == 2: - header += " ratio |" - click.echo(header) - click.echo("|---" * (header.count("|") - 1) + "|") - for workload, num_queries in workloads: - counts = [results[(workload, num_queries, e)][0] for e in engine_names] - row = f"| {workload} | {num_queries} | " - row += " | ".join(f"{c / 1e9:.3f} G" for c in counts) + " |" - if len(engine_names) == 2: - row += f" {counts[1] / counts[0]:.3f} |" - click.echo(row) - paths = {results[(workload, num_queries, e)][1] for e in engine_names} - if len(paths) > 1: + futures = [executor.submit(callgrind, data_dir, *task) for task in tasks] + for future in futures: + future.result() + + click.echo("| workload | queries | tsinfer | tsls | ratio |") + click.echo("|---|---|---|---|---|") + for workload, num_queries, max_sites in workloads: + counts = [] + paths = [] + for engine_name, directory in [("tsinfer", baseline_dir), ("tsls", out_dir)]: + path = out_path(directory, workload, num_queries, max_sites, engine_name) + counts.append(read_total_ir(path)) + paths.append(path.with_suffix(".paths").read_text()) + label = ( + f"{num_queries}" + if max_sites is None + else f"{num_queries} ({max_sites} sites)" + ) + click.echo( + f"| {workload} | {label} | {counts[0] / 1e9:.3f} G | " + f"{counts[1] / 1e9:.3f} G | {counts[1] / counts[0]:.3f} |" + ) + if paths[0] != paths[1]: click.echo(f"WARNING: engines found different paths for {workload}") From 4afaf769ddb117d6543a5af8403dcfd49560fd18 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:11:35 +0100 Subject: [PATCH 06/25] Reset Viterbi per-site state only over the matched window find_path_matrix reset the traceback and max_likelihood_node arrays over every site of the sequence after each call. Each holds an entry per site (581k for chr20p, 16 MB together), so for a short tsinfer ancestor the resets cost far more than the matching: on g84 (1,390 sites per ancestor) memset was 71% of all instructions. Only sites in [start, end) can have been set, so only these are reset now, including when the run fails. Instruction counts (find_path only), before -> after: | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.486 G | 0.162 G | 0.547 | | ancestors_g1000 | 20 | 0.621 G | 0.686 G | 0.362 G | 0.583 | | ancestors_g1483 | 10 | 1.355 G | 1.381 G | 1.221 G | 0.901 | | ancestors_g1484 | 3 | 0.646 G | 0.736 G | 0.688 G | 1.065 | | ancestors_g1485 | 5 | 1.131 G | 1.117 G | 1.036 G | 0.916 | | samples | 1 (100000 sites) | 2.064 G | 2.578 G | 2.565 G | 1.242 | Callgrind counts each iteration of memset's inner loop, so the wall time saving is likely smaller than these counts suggest, but the resets did touch 16 MB per query. Paths unchanged. --- tsls/lib/hmm.c | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index 73f8336..8511e2c 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -1558,6 +1558,7 @@ tsls_hmm_find_path_matrix(tsls_hmm_t *self, tsk_id_t start, tsk_id_t *path_parent, compressed_matrix_t *output) { int ret = 0; + bool valid_window = false; if (start < 0 || end > (tsk_id_t) self->num_sites || start >= end) { ret = TSLS_ERR_BAD_SITE_RANGE; @@ -1567,6 +1568,7 @@ tsls_hmm_find_path_matrix(tsls_hmm_t *self, tsk_id_t start, ret = TSLS_ERR_BAD_SITE_RANGE; goto out; } + valid_window = true; ret = tsls_hmm_reset_traceback(self); if (ret != 0) { goto out; @@ -1578,10 +1580,15 @@ tsls_hmm_find_path_matrix(tsls_hmm_t *self, tsk_id_t start, ret = tsls_hmm_run_traceback(self, start, end, matched_haplotype, path_length, path_left, path_right, path_parent); out: - /* Reset some memory for the next call */ - memset(self->traceback, 0, self->num_sites * sizeof(*self->traceback)); - memset(self->max_likelihood_node, 0xff, - self->num_sites * sizeof(*self->max_likelihood_node)); + /* Reset the per-site state for the next call. Only the sites in the + * window can have been set, and resetting just these matters for short + * windows (tsinfer's early ancestors) on long sequences. */ + if (valid_window) { + memset(self->traceback + start, 0, + (size_t) (end - start) * sizeof(*self->traceback)); + memset(self->max_likelihood_node + start, 0xff, + (size_t) (end - start) * sizeof(*self->max_likelihood_node)); + } memset(self->allelic_state, 0xff, self->num_nodes * sizeof(*self->allelic_state)); return ret; } From 5af3c2e5e0101fb9e0121f36f4ee21314479d02b Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:13:44 +0100 Subject: [PATCH 07/25] Restore the root-marked allele lookup in the site loops Phase 1 dropped this special case from bafe726 to keep the refactor simple. The profile showed the root check in get_allelic_state's walk up the tree costing 0.57G of the sample's 2.57G, more than the whole gap to tsinfer-hmm, so it earns its place: every root's allele is set to the ancestral state from the likelihood list (roots always have likelihoods) before the site's update, and the three likelihood loops walk with get_marked_allelic_state, which has no root check and no assert. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.162 G | 0.159 G | 0.537 | | ancestors_g1000 | 20 | 0.621 G | 0.362 G | 0.340 G | 0.547 | | ancestors_g1483 | 10 | 1.355 G | 1.221 G | 1.171 G | 0.864 | | ancestors_g1484 | 3 | 0.646 G | 0.688 G | 0.608 G | 0.941 | | ancestors_g1485 | 5 | 1.131 G | 1.036 G | 0.963 G | 0.852 | | samples | 1 (100000 sites) | 2.064 G | 2.565 G | 2.232 G | 1.081 | Paths unchanged. --- tsls/lib/hmm.c | 74 ++++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 65 insertions(+), 9 deletions(-) diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index 8511e2c..4b15b94 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -95,6 +95,21 @@ 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 tsls_hmm_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 HMM, so there is deliberately no assert here (the Python module is + * compiled with asserts on); the roots are checked in tsls_hmm_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]; +} + /* Position the specified null tree position at the tree containing x, * seeking from whichever end of the sequence is closer, as * tsk_tree_seek_from_null does. The edges of the tree are then the edges @@ -874,6 +889,47 @@ tsls_hmm_unset_allelic_state(tsls_hmm_t *self, const tsk_id_t site) } } +/* 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 tsls_hmm_set_allelic_state, so that + * mutations above a root take precedence, and the likelihood list must not + * change before tsls_hmm_unset_root_allelic_state. */ +static inline void +tsls_hmm_set_root_allelic_state(tsls_hmm_t *self, const tsk_id_t site) +{ + const allele_t ancestral_state = self->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 +tsls_hmm_unset_root_allelic_state(tsls_hmm_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 @@ -1146,12 +1202,12 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site, const double mu = self->mismatch_rate[site]; const double n = tsls_hmm_get_n(self); const double num_alleles = (double) self->num_alleles[site]; - const allele_t ancestral_state = self->ancestral_state[site]; double max_L, p_no_recomb, p_recomb, p_t, p_e; tsk_id_t u, max_L_node; int j; tsls_hmm_insert_mutation_likelihoods(self, site); + tsls_hmm_set_root_allelic_state(self, site); tsls_hmm_set_allelic_state(self, site); max_L = -1; max_L_node = NULL_NODE; @@ -1167,8 +1223,7 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site, 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]; @@ -1176,6 +1231,7 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site, } } tsls_hmm_unset_allelic_state(self, site); + tsls_hmm_unset_root_allelic_state(self); (void) num_likelihood_nodes; if (max_L <= 0) { @@ -1232,7 +1288,6 @@ tsls_hmm_process_site_forward(tsls_hmm_t *self, const tsk_id_t site, const double mu = self->mismatch_rate[site]; const double n = tsls_hmm_get_n(self); const double num_alleles = (double) self->num_alleles[site]; - const allele_t ancestral_state = self->ancestral_state[site]; double total, normalisation_factor, p_e; tsk_id_t u; int j; @@ -1243,15 +1298,16 @@ tsls_hmm_process_site_forward(tsls_hmm_t *self, const tsk_id_t site, if (total < 0) { total = tsls_hmm_total_likelihood(self); } + tsls_hmm_set_root_allelic_state(self, site); tsls_hmm_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; } tsls_hmm_unset_allelic_state(self, site); + tsls_hmm_unset_root_allelic_state(self); normalisation_factor = tsls_hmm_total_likelihood(self); if (!(normalisation_factor > 0)) { ret = TSLS_ERR_MATCH_IMPOSSIBLE; @@ -1297,7 +1353,6 @@ tsls_hmm_process_site_backward(tsls_hmm_t *self, const tsk_id_t site, const double mu = self->mismatch_rate[site]; const double n = tsls_hmm_get_n(self); const double num_alleles = (double) self->num_alleles[site]; - const allele_t ancestral_state = self->ancestral_state[site]; double x; tsk_id_t u; int j; @@ -1318,14 +1373,15 @@ tsls_hmm_process_site_backward(tsls_hmm_t *self, const tsk_id_t site, goto out; } tsls_hmm_compute_num_inheriting(self); + tsls_hmm_set_root_allelic_state(self, site); tsls_hmm_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); } tsls_hmm_unset_allelic_state(self, site); + tsls_hmm_unset_root_allelic_state(self); transition->pending = true; transition->total = tsls_hmm_total_likelihood(self); transition->rho = self->recombination_rate[site]; From 35c2656201ad82119ae1d437ebc6db528b766a10 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:16:50 +0100 Subject: [PATCH 08/25] Remove asserts from the inner walks of the likelihood cache The walks up the tree in get_inherited_likelihood, its _and_flag variant and insert_mutation_likelihoods had an assert at every step. The Python module builds with asserts on, and tsinfer-hmm's equivalent walks have none. The invariant (every root has a likelihood) is still checked by tsls_hmm_check_state under extended checks, and the once-per-call value asserts remain. The NULL_NODE test in reset_likelihood_cache's walk looks similar but is needed: within one tree update a node with a cached value can become a root (when a later edge in the batch is removed), so the walk can reach the top of the tree. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.159 G | 0.158 G | 0.534 | | ancestors_g1000 | 20 | 0.621 G | 0.340 G | 0.337 G | 0.543 | | ancestors_g1483 | 10 | 1.355 G | 1.171 G | 1.153 G | 0.851 | | ancestors_g1484 | 3 | 0.646 G | 0.608 G | 0.602 G | 0.932 | | ancestors_g1485 | 5 | 1.131 G | 0.963 G | 0.955 G | 0.845 | | samples | 1 (100000 sites) | 2.064 G | 2.232 G | 2.200 G | 1.066 | Paths unchanged. --- tsls/lib/hmm.c | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index 4b15b94..180c447 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -680,7 +680,10 @@ tsls_hmm_update_num_below(tsls_hmm_t *self, tsk_id_t u, int64_t delta) } /* Returns the likelihood inherited by node u, filling in the likelihood - * cache along the path so that later queries are fast. */ + * cache along the path so that later queries are fast. The walks up the + * tree here and in tsls_hmm_get_inherited_likelihood_and_flag are inner + * loops, so they have no asserts (every root has a likelihood, which + * tsls_hmm_check_state checks). */ static double tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u) { @@ -693,7 +696,6 @@ tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u) v = u; while (likely(L[v] == NULL_LIKELIHOOD) && likely(L_cache[v] == CACHE_UNSET)) { v = parent[v]; - assert(v != NULL_NODE); } value = L_cache[v]; if (unlikely(value == CACHE_UNSET)) { @@ -948,7 +950,6 @@ tsls_hmm_insert_mutation_likelihoods(tsls_hmm_t *self, const tsk_id_t site) u = node; while (L[u] == NULL_LIKELIHOOD) { u = parent[u]; - assert(u != NULL_NODE); } tsls_hmm_add_likelihood(self, node, L[u]); } @@ -974,7 +975,6 @@ tsls_hmm_get_inherited_likelihood_and_flag( 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]; From 048bcd3b5a8fd45411d63ad5ff28ab27161e9fdf Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:21:30 +0100 Subject: [PATCH 09/25] Profile the first 100k sites of the sample, not the middle The profiling sample was restricted to the middle 100,000 of its sites, which made both engines build the first tree half way along the sequence. A whole sample (as tsinfer matches it) starts at the first site, so this measured seek costs that the real workload doesn't have: about 100M of tsls's 2.20G. The window is now the first 100,000 sites. tsinfer-hmm's counts are remeasured (its cache files are keyed by the window). | window | tsinfer | tsls | ratio | |---|---|---|---| | middle 100k sites | 2.064 G | 2.200 G | 1.066 | | first 100k sites | 1.642 G | 1.736 G | 1.057 | Most of the gap remains, so it is not mainly the seek. Parity will also be checked on a whole sample once at the end. --- benchmarking/callgrind.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/benchmarking/callgrind.py b/benchmarking/callgrind.py index 3e208e1..061a98b 100644 --- a/benchmarking/callgrind.py +++ b/benchmarking/callgrind.py @@ -36,7 +36,8 @@ # (workload, number of queries, maximum number of sites per query). Each is # about 2G instructions in tsinfer-hmm. A whole sample is about 10G, so the -# sample is restricted to the middle of its window. +# sample is restricted to the start of its window: like a whole sample, its +# forward sweep then starts at the first site, without a long seek. WORKLOADS = [ ("ancestors_g84", 20, None), ("ancestors_g1000", 20, None), @@ -59,12 +60,11 @@ def make_engine(name, ts, model): def clip_window(start, end, max_sites): """ - Return the middle max_sites sites of the window [start, end). + Return the first max_sites sites of the window [start, end). """ - if max_sites is None or end - start <= max_sites: + if max_sites is None: return start, end - start = start + (end - start - max_sites) // 2 - return start, start + max_sites + return start, min(end, start + max_sites) @click.group() @@ -85,7 +85,7 @@ def cli(): def run(workload_dir, engine_name, num_queries, max_sites, instrument): """ Match the first NUM_QUERIES queries in WORKLOAD_DIR with the specified - engine, each over at most MAX_SITES sites in the middle of its window, + engine, each over at most MAX_SITES sites at the start of its window, printing the parents of each path. """ workload_dir = pathlib.Path(workload_dir) @@ -115,7 +115,7 @@ def read_total_ir(path): def out_path(out_dir, workload, num_queries, max_sites, engine_name): - sites = "" if max_sites is None else f".s{max_sites}" + sites = "" if max_sites is None else f".first{max_sites}" return pathlib.Path(out_dir) / ( f"callgrind.{workload}.n{num_queries}{sites}.{engine_name}.out" ) From b0b5a63e5c91dc6a878d0f14332ac80793f3f8b5 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:21:30 +0100 Subject: [PATCH 10/25] Inline the per-edge likelihood helpers insert_edge was called out of line 778k times on the sample (36M instructions), and get_inherited_likelihood 770k times, where tsinfer-hmm does the same work inline in its tree loop. Marking these and remove_edge/reset_likelihood_cache inline saves a little. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.158 G | 0.158 G | 0.533 | | ancestors_g1000 | 20 | 0.621 G | 0.337 G | 0.336 G | 0.541 | | ancestors_g1483 | 10 | 1.355 G | 1.153 G | 1.147 G | 0.846 | | ancestors_g1484 | 3 | 0.646 G | 0.602 G | 0.600 G | 0.928 | | ancestors_g1485 | 5 | 1.131 G | 0.955 G | 0.951 G | 0.841 | | samples | 1 (first 100k sites) | 1.642 G | 1.736 G | 1.716 G | 1.045 | Investigated and not the cause: the cache-reset walk's NULL_NODE test. Both engines spend the same on cache resets (tsls 107M; tsinfer-hmm 80M + 24M), so making that walk cheaper would not close the gap. Paths unchanged. --- tsls/lib/hmm.c | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index 180c447..55de2ac 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -684,7 +684,7 @@ tsls_hmm_update_num_below(tsls_hmm_t *self, tsk_id_t u, int64_t delta) * tree here and in tsls_hmm_get_inherited_likelihood_and_flag are inner * loops, so they have no asserts (every root has a likelihood, which * tsls_hmm_check_state checks). */ -static double +static inline double tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u) { const double *restrict L = self->likelihood; @@ -710,7 +710,7 @@ tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u) return value; } -static void +static inline void tsls_hmm_reset_likelihood_cache(tsls_hmm_t *self, tsk_id_t u) { double *restrict L_cache = self->likelihood_cache; @@ -727,7 +727,7 @@ tsls_hmm_reset_likelihood_cache(tsls_hmm_t *self, tsk_id_t u) * isolated after removing the edge are added to the isolated list; they * only leave the tree if they are still isolated after all the edges at * this position have been inserted (see tsls_hmm_update_tree). */ -static void +static inline void tsls_hmm_remove_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, bool update_likelihoods, tsk_id_t *restrict isolated, int *num_isolated) { @@ -759,7 +759,7 @@ tsls_hmm_remove_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, /* Insert the edge from p to c into the tree. Nodes entering the tree get * likelihood 0, and so can only be reached by recombination. */ -static void +static inline void tsls_hmm_insert_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, bool update_likelihoods) { double *restrict L = self->likelihood; From 4c565478ae85ae8f84ee156c28e58bf8b9d52ceb Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:29:29 +0100 Subject: [PATCH 11/25] Keep the roots marked through the sweep instead of at each site Each site marked every root's allele before its update and unmarked them after, with two passes over the likelihood list (38M of the sample's 1.72G; tsinfer-hmm has a single root and no such passes). Now every node without a parent is permanently marked with ROOT_ALLELE during the sweeps: init_sweep sets the marks with a memset, insert_edge and remove_edge update the child's, and unsetting a mutation restores the mark if the node is a root. get_marked_allelic_state turns ROOT_ALLELE into the site's ancestral state with one compare per walk. The traceback, which doesn't maintain marks, starts from a cleared array. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.158 G | 0.161 G | 0.544 | | ancestors_g1000 | 20 | 0.621 G | 0.336 G | 0.339 G | 0.547 | | ancestors_g1483 | 10 | 1.355 G | 1.147 G | 1.157 G | 0.854 | | ancestors_g1484 | 3 | 0.646 G | 0.600 G | 0.599 G | 0.928 | | ancestors_g1485 | 5 | 1.131 G | 0.951 G | 0.963 G | 0.852 | | samples | 1 (first 100k sites) | 1.642 G | 1.716 G | 1.691 G | 1.030 | This trades about 1% on ancestors (marking the first tree's edges and clearing for the traceback) for 1.5% on the sample, the one workload still above tsinfer-hmm. A first version set the marks in init_sweep's per-node loop, which cost 2-4% on ancestors (0.973G on g1485); a memset fixed most of that. Dead end, reverted: marking the traceback tree's roots from the site's traceback list (which contains every root), so the traceback's allele lookup could also skip the root check. The two extra passes over the list cost more than the check: the sample went from 1.716G to 1.734G, and every ancestor workload got worse too. Paths unchanged. --- tsls/lib/hmm.c | 101 ++++++++++++++++++++++--------------------------- 1 file changed, 46 insertions(+), 55 deletions(-) diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index 55de2ac..f86fe42 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -95,19 +95,27 @@ 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 tsls_hmm_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 HMM, so there is deliberately no assert here (the Python module is - * compiled with asserts on); the roots are checked in tsls_hmm_check_state. */ +/* During the forward and backward sweeps, the allelic state of every node + * without a parent is ROOT_ALLELE, unless it has a mutation at the current + * site. insert_edge and remove_edge keep this up to date. */ +#define ROOT_ALLELE (-2) + +/* As get_allelic_state, but for a tree in which the roots are marked with + * ROOT_ALLELE, so the walk up the tree doesn't need to check for the root. + * This is the innermost loop of the HMM, so there is deliberately no assert + * here (the Python module is compiled with asserts on); the marks are + * checked in tsls_hmm_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) +get_marked_allelic_state(const allele_t *restrict allelic_state, + const tsk_id_t *restrict parent, tsk_id_t u, allele_t ancestral_state) { + allele_t state; + while (allelic_state[u] == TSK_NULL) { u = parent[u]; } - return allelic_state[u]; + state = allelic_state[u]; + return likely(state != ROOT_ALLELE) ? state : ancestral_state; } /* Position the specified null tree position at the tree containing x, @@ -279,7 +287,8 @@ tsls_hmm_check_state(const tsls_hmm_t *self) } else { assert(self->likelihood[u] == NOT_IN_TREE); } - assert(self->allelic_state[u] == TSK_NULL); + assert(self->allelic_state[u] + == (self->parent[u] == NULL_NODE ? ROOT_ALLELE : TSK_NULL)); assert(self->likelihood_cache[u] == CACHE_UNSET); } assert(num_likelihoods == self->num_likelihood_nodes); @@ -741,6 +750,7 @@ tsls_hmm_remove_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, } remove_edge(p, c, self->parent, self->left_child, self->right_child, self->left_sib, self->right_sib); + self->allelic_state[c] = ROOT_ALLELE; if (update_likelihoods) { /* Each node can only become isolated once when removing edges, so * the isolated list has at most num_nodes entries. */ @@ -776,6 +786,7 @@ tsls_hmm_insert_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, bool update_likel } insert_edge(p, c, self->parent, self->left_child, self->right_child, self->left_sib, self->right_sib); + self->allelic_state[c] = TSK_NULL; if (self->track_num_below) { tsls_hmm_update_num_below(self, p, self->num_below[c]); } @@ -839,6 +850,8 @@ tsls_hmm_init_sweep( memset(self->left_sib, 0xff, N * sizeof(*self->left_sib)); memset(self->right_sib, 0xff, N * sizeof(*self->right_sib)); memset(self->recombination_required, 0xff, N * sizeof(*self->recombination_required)); + /* Every node starts as a root (ROOT_ALLELE is 0xfe as an int8) */ + memset(self->allelic_state, (uint8_t) ROOT_ALLELE, N * sizeof(*self->allelic_state)); for (k = 0; k < N; k++) { self->likelihood_cache[k] = CACHE_UNSET; self->num_below[k] = 1; @@ -879,56 +892,33 @@ tsls_hmm_set_allelic_state(tsls_hmm_t *self, const tsk_id_t site) } } +/* Unset the allelic states set for the site's mutations, keeping the + * roots marked with ROOT_ALLELE (during the sweeps). */ static inline void tsls_hmm_unset_allelic_state(tsls_hmm_t *self, const tsk_id_t site) { const tsk_mutation_t *restrict mutations = self->tree_sequence->site_mutations[site]; const tsk_size_t num_mutations = self->tree_sequence->site_mutations_length[site]; tsk_size_t j; - - for (j = 0; j < num_mutations; j++) { - self->allelic_state[mutations[j].node] = TSK_NULL; - } -} - -/* 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 tsls_hmm_set_allelic_state, so that - * mutations above a root take precedence, and the likelihood list must not - * change before tsls_hmm_unset_root_allelic_state. */ -static inline void -tsls_hmm_set_root_allelic_state(tsls_hmm_t *self, const tsk_id_t site) -{ - const allele_t ancestral_state = self->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; - } + for (j = 0; j < num_mutations; j++) { + u = mutations[j].node; + self->allelic_state[u] = self->parent[u] == NULL_NODE ? ROOT_ALLELE : TSK_NULL; } } +/* As tsls_hmm_unset_allelic_state, for the traceback, which doesn't mark + * the roots. */ static inline void -tsls_hmm_unset_root_allelic_state(tsls_hmm_t *self) +tsls_hmm_clear_allelic_state(tsls_hmm_t *self, const tsk_id_t 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; + const tsk_mutation_t *restrict mutations = self->tree_sequence->site_mutations[site]; + const tsk_size_t num_mutations = self->tree_sequence->site_mutations_length[site]; + tsk_size_t j; - for (j = 0; j < self->num_likelihood_nodes; j++) { - u = L_nodes[j]; - if (parent[u] == NULL_NODE) { - allelic_state[u] = TSK_NULL; - } + for (j = 0; j < num_mutations; j++) { + self->allelic_state[mutations[j].node] = TSK_NULL; } } @@ -1202,12 +1192,12 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site, const double mu = self->mismatch_rate[site]; const double n = tsls_hmm_get_n(self); const double num_alleles = (double) self->num_alleles[site]; + const allele_t ancestral_state = self->ancestral_state[site]; double max_L, p_no_recomb, p_recomb, p_t, p_e; tsk_id_t u, max_L_node; int j; tsls_hmm_insert_mutation_likelihoods(self, site); - tsls_hmm_set_root_allelic_state(self, site); tsls_hmm_set_allelic_state(self, site); max_L = -1; max_L_node = NULL_NODE; @@ -1223,7 +1213,8 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site, recombination_required[u] = true; } p_e = emission_probability( - get_marked_allelic_state(allelic_state, parent, u), state, mu, num_alleles); + get_marked_allelic_state(allelic_state, parent, u, ancestral_state), state, + mu, num_alleles); L[u] = p_t * p_e; if (L[u] > max_L) { max_L = L[u]; @@ -1231,7 +1222,6 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site, } } tsls_hmm_unset_allelic_state(self, site); - tsls_hmm_unset_root_allelic_state(self); (void) num_likelihood_nodes; if (max_L <= 0) { @@ -1288,6 +1278,7 @@ tsls_hmm_process_site_forward(tsls_hmm_t *self, const tsk_id_t site, const double mu = self->mismatch_rate[site]; const double n = tsls_hmm_get_n(self); const double num_alleles = (double) self->num_alleles[site]; + const allele_t ancestral_state = self->ancestral_state[site]; double total, normalisation_factor, p_e; tsk_id_t u; int j; @@ -1298,16 +1289,15 @@ tsls_hmm_process_site_forward(tsls_hmm_t *self, const tsk_id_t site, if (total < 0) { total = tsls_hmm_total_likelihood(self); } - tsls_hmm_set_root_allelic_state(self, site); tsls_hmm_set_allelic_state(self, site); for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; p_e = emission_probability( - get_marked_allelic_state(allelic_state, parent, u), state, mu, num_alleles); + get_marked_allelic_state(allelic_state, parent, u, ancestral_state), state, + mu, num_alleles); L[u] = (L[u] * (1 - rho) + rho / n * total) * p_e; } tsls_hmm_unset_allelic_state(self, site); - tsls_hmm_unset_root_allelic_state(self); normalisation_factor = tsls_hmm_total_likelihood(self); if (!(normalisation_factor > 0)) { ret = TSLS_ERR_MATCH_IMPOSSIBLE; @@ -1353,6 +1343,7 @@ tsls_hmm_process_site_backward(tsls_hmm_t *self, const tsk_id_t site, const double mu = self->mismatch_rate[site]; const double n = tsls_hmm_get_n(self); const double num_alleles = (double) self->num_alleles[site]; + const allele_t ancestral_state = self->ancestral_state[site]; double x; tsk_id_t u; int j; @@ -1373,15 +1364,14 @@ tsls_hmm_process_site_backward(tsls_hmm_t *self, const tsk_id_t site, goto out; } tsls_hmm_compute_num_inheriting(self); - tsls_hmm_set_root_allelic_state(self, site); tsls_hmm_set_allelic_state(self, site); for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; L[u] *= emission_probability( - get_marked_allelic_state(allelic_state, parent, u), state, mu, num_alleles); + get_marked_allelic_state(allelic_state, parent, u, ancestral_state), state, + mu, num_alleles); } tsls_hmm_unset_allelic_state(self, site); - tsls_hmm_unset_root_allelic_state(self); transition->pending = true; transition->total = tsls_hmm_total_likelihood(self); transition->rho = self->recombination_rate[site]; @@ -1522,6 +1512,7 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, memset(parent, 0xff, self->num_nodes * sizeof(*parent)); memset( recombination_required, 0xff, self->num_nodes * sizeof(*recombination_required)); + memset(self->allelic_state, 0xff, self->num_nodes * sizeof(*self->allelic_state)); x = sites_position[end - 1]; ret = seek_tree_position(&tp, x); if (ret != 0) { @@ -1550,7 +1541,7 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, u = path_parent[path_length]; match[site] = get_allelic_state(allelic_state, parent, u, self->ancestral_state[site]); - tsls_hmm_unset_allelic_state(self, site); + tsls_hmm_clear_allelic_state(self, site); /* Mark the traceback nodes on the tree */ tsls_hmm_set_recombination_required(self, site, recombination_required); From c57dd8b1047daa53e818dd2890b2e4fa72c5d484 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:31:55 +0100 Subject: [PATCH 12/25] Force inlining of the coalescing helpers The walks in coalesce_likelihoods cost the same as tsinfer-hmm's (133M + 106M against 134M + 100M on the sample), but tsls spent about 40M more around them: get_inherited_likelihood_and_flag was not inlined despite 'inline' (17M of per-call setup), each call asserted value >= 0 (5.8M), and the loop tested the invariant check_recombination_required for every node (5.8M). The helper and coalesce_likelihoods are now always_inline, so each call site gets a copy specialised for its constant flag, and the per-call value asserts are gone (as for the walks in the last change). The flag cache stores (10M) remain; they are the cost of coalescing with recombination flags, which halves the traceback. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.161 G | 0.160 G | 0.542 | | ancestors_g1000 | 20 | 0.621 G | 0.339 G | 0.338 G | 0.545 | | ancestors_g1483 | 10 | 1.355 G | 1.157 G | 1.153 G | 0.851 | | ancestors_g1484 | 3 | 0.646 G | 0.599 G | 0.597 G | 0.924 | | ancestors_g1485 | 5 | 1.131 G | 0.963 G | 0.961 G | 0.849 | | samples | 1 (first 100k sites) | 1.642 G | 1.691 G | 1.680 G | 1.023 | Paths unchanged. --- tsls/lib/hmm.c | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index f86fe42..04f59c8 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -33,6 +33,15 @@ #include #include +/* For the innermost helpers, which GCC doesn't always inline. Inlining + * tsls_hmm_coalesce_likelihoods also specialises it for its constant + * check_recombination_required argument. */ +#ifdef __GNUC__ +#define ALWAYS_INLINE inline __attribute__((always_inline)) +#else +#define ALWAYS_INLINE inline +#endif + #define MODE_VITERBI 0 #define MODE_FORWARD 1 #define MODE_BACKWARD 2 @@ -710,7 +719,6 @@ tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u) if (unlikely(value == CACHE_UNSET)) { value = L[v]; } - assert(value >= 0); v = u; while (likely(L[v] == NULL_LIKELIHOOD) && likely(L_cache[v] == CACHE_UNSET)) { L_cache[v] = value; @@ -949,7 +957,7 @@ tsls_hmm_insert_mutation_likelihoods(tsls_hmm_t *self, const tsk_id_t site) /* As tsls_hmm_get_inherited_likelihood, but also return the * recombination_required flag that u inherits, caching it along the path * with the likelihood. */ -static inline double +static ALWAYS_INLINE double tsls_hmm_get_inherited_likelihood_and_flag( tsls_hmm_t *self, tsk_id_t u, int8_t *flag) { @@ -973,7 +981,6 @@ tsls_hmm_get_inherited_likelihood_and_flag( 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; @@ -989,7 +996,7 @@ tsls_hmm_get_inherited_likelihood_and_flag( * 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 +static ALWAYS_INLINE void tsls_hmm_coalesce_likelihoods( tsls_hmm_t *self, bool check_recombination_required) { From ebdc10f0d2177df9090fb8af574cbca96b5bef7c Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:34:25 +0100 Subject: [PATCH 13/25] Give the traceback tree a virtual root The traceback looked up the matched allele with get_allelic_state, which checks for a root at every step of the walk (34M on the sample, against 23M for tsinfer-hmm's walk to its single root). Now every root of the traceback tree has the virtual root N as its parent (the tree arrays already have N + 1 entries), and the virtual root's allelic state is set to the site's ancestral state, so the walk needs no root check. Setting up the parent array is one fill per traceback instead of a memset. get_allelic_state is no longer used and is removed. This replaces the dead end in 4c56547 (marking the traceback roots from the site's traceback list), which cost two passes over the list per site. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.160 G | 0.156 G | 0.529 | | ancestors_g1000 | 20 | 0.621 G | 0.338 G | 0.328 G | 0.528 | | ancestors_g1483 | 10 | 1.355 G | 1.153 G | 1.133 G | 0.836 | | ancestors_g1484 | 3 | 0.646 G | 0.597 G | 0.591 G | 0.915 | | ancestors_g1485 | 5 | 1.131 G | 0.961 G | 0.950 G | 0.840 | | samples | 1 (first 100k sites) | 1.642 G | 1.680 G | 1.670 G | 1.017 | Paths unchanged. --- tsls/lib/hmm.c | 51 ++++++++++++++++++++++++-------------------------- 1 file changed, 24 insertions(+), 27 deletions(-) diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index 04f59c8..d1a4737 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -89,31 +89,17 @@ insert_edge(tsk_id_t p, tsk_id_t c, tsk_id_t *restrict parent, right_child[p] = c; } -/* Returns the allelic state of node u, given the allelic_state array in - * which the nodes with mutations at the current site are set. */ -static inline allele_t -get_allelic_state(const allele_t *restrict allelic_state, - const tsk_id_t *restrict parent, tsk_id_t u, allele_t ancestral_state) -{ - while (allelic_state[u] == TSK_NULL) { - if (parent[u] == NULL_NODE) { - return ancestral_state; - } - u = parent[u]; - } - return allelic_state[u]; -} - /* During the forward and backward sweeps, the allelic state of every node * without a parent is ROOT_ALLELE, unless it has a mutation at the current * site. insert_edge and remove_edge keep this up to date. */ #define ROOT_ALLELE (-2) -/* As get_allelic_state, but for a tree in which the roots are marked with - * ROOT_ALLELE, so the walk up the tree doesn't need to check for the root. - * This is the innermost loop of the HMM, so there is deliberately no assert - * here (the Python module is compiled with asserts on); the marks are - * checked in tsls_hmm_check_state. */ +/* Returns the allelic state of node u, given the allelic_state array in + * which the nodes with mutations at the current site are set and the roots + * are marked with ROOT_ALLELE, so the walk up the tree doesn't need to + * check for the root. This is the innermost loop of the HMM, so there is + * deliberately no assert here (the Python module is compiled with asserts + * on); the marks are checked in tsls_hmm_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, allele_t ancestral_state) @@ -1492,9 +1478,13 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, double x; bool valid; tsk_id_t *restrict parent = self->parent; - const allele_t *restrict allelic_state = self->allelic_state; + allele_t *restrict allelic_state = self->allelic_state; int8_t *restrict recombination_required = self->recombination_required; const double *restrict sites_position = self->sites_position; + /* Every root of the traceback tree has the virtual root as its parent, + * whose allelic state is the site's ancestral state, so the walk to find + * a node's allele doesn't need to check for the root. */ + const tsk_id_t virtual_root = (tsk_id_t) self->num_nodes; const tsk_edge_table_t *edges = &self->tree_sequence->tables->edges; const tsk_id_t *restrict edge_parent = edges->parent; const tsk_id_t *restrict edge_child = edges->child; @@ -1516,10 +1506,13 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, /* Build the tree at the last site, and then go through the trees in * reverse, running the traceback at each site. */ - memset(parent, 0xff, self->num_nodes * sizeof(*parent)); + for (u = 0; u < virtual_root; u++) { + parent[u] = virtual_root; + } + parent[virtual_root] = NULL_NODE; memset( recombination_required, 0xff, self->num_nodes * sizeof(*recombination_required)); - memset(self->allelic_state, 0xff, self->num_nodes * sizeof(*self->allelic_state)); + memset(allelic_state, 0xff, self->num_nodes * sizeof(*allelic_state)); x = sites_position[end - 1]; ret = seek_tree_position(&tp, x); if (ret != 0) { @@ -1537,18 +1530,22 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, assert(valid); (void) valid; for (j = tp.out.start; j != tp.out.stop; j--) { - parent[edge_child[tp.out.order[j]]] = NULL_NODE; + parent[edge_child[tp.out.order[j]]] = virtual_root; } for (j = tp.in.start; j != tp.in.stop; j--) { e = tp.in.order[j]; parent[edge_child[e]] = edge_parent[e]; } } + allelic_state[virtual_root] = self->ancestral_state[site]; tsls_hmm_set_allelic_state(self, site); u = path_parent[path_length]; - match[site] - = get_allelic_state(allelic_state, parent, u, self->ancestral_state[site]); + while (allelic_state[u] == TSK_NULL) { + u = parent[u]; + } + match[site] = allelic_state[u]; tsls_hmm_clear_allelic_state(self, site); + u = path_parent[path_length]; /* Mark the traceback nodes on the tree */ tsls_hmm_set_recombination_required(self, site, recombination_required); @@ -1558,7 +1555,7 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, * root has a likelihood, and so is marked. */ while (recombination_required[u] == -1) { u = parent[u]; - assert(u != NULL_NODE); + assert(u != virtual_root); } if (recombination_required[u] && site > start) { max_likelihood_node = self->max_likelihood_node[site - 1]; From 0cd8ab35351441812374cc096f06e4d4f94cbca2 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:39:21 +0100 Subject: [PATCH 14/25] Drop the root test from the cache reset walk reset_likelihood_cache tested u != NULL_NODE at every step, because within one tree update a node with a cached value could become a root: when an edge above a cached node is removed later in the same batch. Now remove_edge clears the cut child's cached value. That value is redundant, since the child is given the same value as its likelihood. So no root ever has a cached value. Values are only cached on nodes without likelihoods, and every root has one, so every cached chain ends below a root and the walk can stop at the first unset value. This is the same loop as tsinfer-hmm's (which relies on its single root never being cached). Checked under valgrind memcheck with the dense comparisons, window, Viterbi matrix and pruning tests (138 tests with extended checks, which also verify that every cached value is unset after each update): no invalid reads, so the walks never leave the tree. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.156 G | 0.155 G | 0.523 | | ancestors_g1000 | 20 | 0.621 G | 0.328 G | 0.323 G | 0.520 | | ancestors_g1483 | 10 | 1.355 G | 1.133 G | 1.110 G | 0.819 | | ancestors_g1484 | 3 | 0.646 G | 0.591 G | 0.583 G | 0.902 | | ancestors_g1485 | 5 | 1.131 G | 0.950 G | 0.939 G | 0.830 | | samples | 1 (first 100k sites) | 1.642 G | 1.670 G | 1.638 G | 0.998 | tsls now uses no more instructions than tsinfer-hmm on every workload of the profiling set. Paths unchanged. --- tsls/lib/hmm.c | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index d1a4737..2d897e8 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -713,13 +713,16 @@ tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u) return value; } +/* Unset the cached values on the path up from u. Values are only cached on + * nodes without likelihoods, and no root has a cached value (see + * tsls_hmm_remove_edge), so the walk stops below a root. */ static inline void tsls_hmm_reset_likelihood_cache(tsls_hmm_t *self, tsk_id_t u) { double *restrict L_cache = self->likelihood_cache; const tsk_id_t *restrict parent = self->parent; - while (u != NULL_NODE && L_cache[u] != CACHE_UNSET) { + while (L_cache[u] != CACHE_UNSET) { L_cache[u] = CACHE_UNSET; u = parent[u]; } @@ -745,6 +748,12 @@ tsls_hmm_remove_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, remove_edge(p, c, self->parent, self->left_child, self->right_child, self->left_sib, self->right_sib); self->allelic_state[c] = ROOT_ALLELE; + /* c is now a root, and so has a likelihood. Its cached value (if any, + * from an earlier removal at this position) is the same as its + * likelihood, and clearing it means that no root has a cached value, + * so walks over the cache always stop below a root (see + * tsls_hmm_reset_likelihood_cache). */ + self->likelihood_cache[c] = CACHE_UNSET; if (update_likelihoods) { /* Each node can only become isolated once when removing edges, so * the isolated list has at most num_nodes entries. */ From ea952ad424a905a1b9f50c6372de163d6cdfa0aa Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:45:06 +0100 Subject: [PATCH 15/25] Profile the sample on a reference cut at 100k sites The previous change matched the first 100k sites of the sample against the whole reference, and showed tsls at parity (0.998). A one-off run on the original full set (20 ancestors per group, 1 whole sample) disagreed for the sample: | workload | tsinfer | tsls | ratio | |---|---|---|---| | ancestors_g84 (20) | 0.296 G | 0.155 G | 0.523 | | ancestors_g1000 (20) | 0.621 G | 0.323 G | 0.520 | | ancestors_g1483 (20) | 2.388 G | 1.877 G | 0.786 | | ancestors_g1484 (20) | 9.911 G | 9.575 G | 0.966 | | ancestors_g1485 (20) | 6.653 G | 5.940 G | 0.893 | | samples (1 whole) | 9.922 G | 10.287 G | 1.037 | The window was biased: tsinfer-hmm's traceback walks in from the end of the sequence, which costs about 60M for a window at the start but nothing extra for a whole sample, while tsls seeks to the window's last site. So the first-100k window overstated tsinfer-hmm's cost by about 4%. (The ancestor windows don't have this problem: real ancestors are matched that way.) Now the sample's reference is cut at site 100,000 (keep_intervals and rtrim, node IDs unchanged) and the sample is matched over the whole of it, so both engines start at the first site and finish at the last, as for a whole sample. This gives 1.561G against 1.618G (1.036), in line with the whole sample's 1.037, so parity has not yet been reached on samples. --- benchmarking/README.md | 5 ++-- benchmarking/callgrind.py | 61 ++++++++++++++++++++++----------------- 2 files changed, 37 insertions(+), 29 deletions(-) diff --git a/benchmarking/README.md b/benchmarking/README.md index 7d7dd68..9920135 100644 --- a/benchmarking/README.md +++ b/benchmarking/README.md @@ -366,9 +366,10 @@ for the whole group divided by the number of haplotypes. 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 + At the time, building tsls's `MatcherIndexes` took several times longer (3.6 s against 0.4 s for the full ancestors tree - sequence), but this is done once per group. + sequence); tsls now uses a `tsls.TreeSequence` (a `tsk_treeseq_t`) + instead, built 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 diff --git a/benchmarking/callgrind.py b/benchmarking/callgrind.py index 061a98b..6d29f43 100644 --- a/benchmarking/callgrind.py +++ b/benchmarking/callgrind.py @@ -34,10 +34,13 @@ "tsls": "tsls_hmm_find_path", } -# (workload, number of queries, maximum number of sites per query). Each is -# about 2G instructions in tsinfer-hmm. A whole sample is about 10G, so the -# sample is restricted to the start of its window: like a whole sample, its -# forward sweep then starts at the first site, without a long seek. +# (workload, number of queries, number of sites). Each is about 2G +# instructions in tsinfer-hmm. A whole sample is about 10G, so for the +# sample the reference is cut at the specified number of sites, and the +# sample is matched over all of them. Matching just the first sites of the +# full reference would not be representative: tsinfer-hmm's traceback walks +# in from the end of the sequence, which a whole sample doesn't need, and +# which made tsinfer-hmm's counts about 4% higher. WORKLOADS = [ ("ancestors_g84", 20, None), ("ancestors_g1000", 20, None), @@ -58,13 +61,13 @@ def make_engine(name, ts, model): ) -def clip_window(start, end, max_sites): +def cut_reference(ts, num_sites): """ - Return the first max_sites sites of the window [start, end). + Return the tree sequence cut just before the specified site, so that it + has only the first num_sites sites. Node IDs are unchanged. """ - if max_sites is None: - return start, end - return start, min(end, start + max_sites) + cut = ts.sites_position[num_sites] + return ts.keep_intervals([[0, cut]], simplify=False).rtrim() @click.group() @@ -76,22 +79,25 @@ def cli(): @click.argument("workload_dir", type=click.Path(exists=True, file_okay=False)) @click.argument("engine_name", type=click.Choice(list(FIND_PATH_FUNCTIONS))) @click.option("-n", "--num-queries", type=int, default=20) -@click.option("--max-sites", type=int, default=None) +@click.option("--num-sites", type=int, default=None) @click.option( "--instrument", is_flag=True, help="Switch on callgrind instrumentation after setup", ) -def run(workload_dir, engine_name, num_queries, max_sites, instrument): +def run(workload_dir, engine_name, num_queries, num_sites, instrument): """ Match the first NUM_QUERIES queries in WORKLOAD_DIR with the specified - engine, each over at most MAX_SITES sites at the start of its window, - printing the parents of each path. + engine, printing the parents of each path. If NUM_SITES is specified, + the reference is cut after that many sites, and the queries are matched + over their windows within these sites. """ workload_dir = pathlib.Path(workload_dir) model = queries.load_model(workload_dir) ts = run_module.load_ts(workload_dir / "reference.trees") H, _, df = queries.load_queries(workload_dir) + if num_sites is not None: + ts = cut_reference(ts, num_sites) engine = make_engine(engine_name, ts, model) if instrument: subprocess.run( @@ -101,8 +107,9 @@ def run(workload_dir, engine_name, num_queries, max_sites, instrument): ) for query_id in df.query_id[:num_queries]: start, end = queries.query_window(H, query_id) - start, end = clip_window(start, end, max_sites) - _, _, parent, _, _ = engine.match(H[query_id], start, end) + h = H[query_id][: ts.num_sites] + end = min(end, ts.num_sites) + _, _, parent, _, _ = engine.match(h, start, end) print(query_id, " ".join(str(p) for p in np.asarray(parent))) @@ -114,14 +121,14 @@ def read_total_ir(path): raise ValueError(f"No totals in {path}") -def out_path(out_dir, workload, num_queries, max_sites, engine_name): - sites = "" if max_sites is None else f".first{max_sites}" +def out_path(out_dir, workload, num_queries, num_sites, engine_name): + sites = "" if num_sites is None else f".cut{num_sites}" return pathlib.Path(out_dir) / ( f"callgrind.{workload}.n{num_queries}{sites}.{engine_name}.out" ) -def callgrind(data_dir, workload, num_queries, max_sites, engine_name, out_file): +def callgrind(data_dir, workload, num_queries, num_sites, engine_name, out_file): cmd = [ "valgrind", "--tool=callgrind", @@ -138,8 +145,8 @@ def callgrind(data_dir, workload, num_queries, max_sites, engine_name, out_file) str(num_queries), "--instrument", ] - if max_sites is not None: - cmd += ["--max-sites", str(max_sites)] + if num_sites is not None: + cmd += ["--num-sites", str(num_sites)] result = subprocess.run(cmd, check=False, capture_output=True, text=True) if result.returncode != 0: raise RuntimeError(f"{' '.join(cmd)} failed:\n{result.stderr[-2000:]}") @@ -179,13 +186,13 @@ def measure(data_dir, workload_names, out_dir, baseline_dir, jobs): if len(workload_names) > 0: workloads = [w for w in WORKLOADS if w[0] in workload_names] tasks = [] - for workload, num_queries, max_sites in workloads: + for workload, num_queries, num_sites in workloads: for engine_name, directory in [("tsinfer", baseline_dir), ("tsls", out_dir)]: pathlib.Path(directory).mkdir(parents=True, exist_ok=True) - path = out_path(directory, workload, num_queries, max_sites, engine_name) + path = out_path(directory, workload, num_queries, num_sites, engine_name) if engine_name == "tsinfer" and path.exists(): continue - tasks.append((workload, num_queries, max_sites, engine_name, path)) + tasks.append((workload, num_queries, num_sites, engine_name, path)) # Run the largest workloads first, so that they don't finish last tasks.sort(key=lambda task: task[0] not in ("samples", "ancestors_g1485")) with concurrent.futures.ThreadPoolExecutor(jobs) as executor: @@ -195,17 +202,17 @@ def measure(data_dir, workload_names, out_dir, baseline_dir, jobs): click.echo("| workload | queries | tsinfer | tsls | ratio |") click.echo("|---|---|---|---|---|") - for workload, num_queries, max_sites in workloads: + for workload, num_queries, num_sites in workloads: counts = [] paths = [] for engine_name, directory in [("tsinfer", baseline_dir), ("tsls", out_dir)]: - path = out_path(directory, workload, num_queries, max_sites, engine_name) + path = out_path(directory, workload, num_queries, num_sites, engine_name) counts.append(read_total_ir(path)) paths.append(path.with_suffix(".paths").read_text()) label = ( f"{num_queries}" - if max_sites is None - else f"{num_queries} ({max_sites} sites)" + if num_sites is None + else f"{num_queries} ({num_sites} sites)" ) click.echo( f"| {workload} | {label} | {counts[0] / 1e9:.3f} G | " From 0736f0d5c3d2dc730d4fb11358b65be326354e1a Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 12:49:22 +0100 Subject: [PATCH 16/25] Remove nodes leaving the tree from the likelihood list in one pass When nodes left the tree at a tree transition, each was removed from the likelihood list with its own pass over the list (delete_likelihood, as in tsinfer-hmm): 96M on the sample, with about 5 nodes leaving per transition. They are now marked NOT_IN_TREE and removed together with a single pass, only when some node has left. | workload | queries | tsinfer | tsls before | tsls after | ratio | |---|---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.155 G | 0.154 G | 0.522 | | ancestors_g1000 | 20 | 0.621 G | 0.323 G | 0.320 G | 0.516 | | ancestors_g1483 | 10 | 1.355 G | 1.110 G | 1.067 G | 0.787 | | ancestors_g1484 | 3 | 0.646 G | 0.583 G | 0.567 G | 0.878 | | ancestors_g1485 | 5 | 1.131 G | 0.939 G | 0.919 G | 0.812 | | samples | 1 (cut at 100k sites) | 1.561 G | 1.618 G | 1.525 G | 0.976 | (The ancestor 'before' numbers are from 0cd8ab3; the sample's is on the cut reference from ea952ad.) Considered and not done: tsls's remaining extra work on the sample is mostly the recombination flag bookkeeping (about 35M: the flag cache stores, the flag compare and the traceback reuse check), which buys the halved traceback, and tskit's tree position scans (about 20M more than tsinfer-hmm's inline edge loops), which come with using tsk_tree_position_t. Packing the flag into the cached value's sign bit would save perhaps 10-15M but needs a new CACHE_UNSET sentinel and care with -0.0; not worth it now. Paths unchanged. --- tsls/lib/hmm.c | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index 2d897e8..5791567 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -654,21 +654,22 @@ tsls_hmm_add_likelihood(tsls_hmm_t *self, tsk_id_t u, double value) self->num_likelihood_nodes++; } -/* Remove the specified node from the list of nodes with likelihoods. */ +/* Remove the nodes that have left the tree (whose likelihood is + * NOT_IN_TREE) from the list of nodes with likelihoods. */ static void -tsls_hmm_delete_likelihood(tsls_hmm_t *self, const tsk_id_t node) +tsls_hmm_remove_left_nodes(tsls_hmm_t *self) { int j, k; tsk_id_t *restrict L_nodes = self->likelihood_nodes; + const double *restrict L = self->likelihood; k = 0; for (j = 0; j < self->num_likelihood_nodes; j++) { L_nodes[k] = L_nodes[j]; - if (L_nodes[j] != node) { + if (L[L_nodes[j]] != NOT_IN_TREE) { k++; } } - assert(self->num_likelihood_nodes == k + 1); self->num_likelihood_nodes = k; } @@ -806,7 +807,7 @@ tsls_hmm_update_tree(tsls_hmm_t *self, const tsk_tree_position_t *tp) tsk_id_t *restrict isolated = self->likelihood_nodes_tmp; double *restrict L = self->likelihood; tsk_id_t j, e, u; - int k, num_isolated = 0; + int k, num_left, num_isolated = 0; for (j = tp->out.start; j != tp->out.stop; j += direction) { e = tp->out.order[j]; @@ -821,14 +822,20 @@ tsls_hmm_update_tree(tsls_hmm_t *self, const tsk_tree_position_t *tp) tsls_hmm_insert_edge(self, edge_parent[e], edge_child[e], true); } /* Nodes that are still isolated have left the tree. Nodes that were - * isolated and then inserted again keep their likelihoods. */ + * isolated and then inserted again keep their likelihoods. We mark the + * nodes that have left and then remove them from the likelihood list in + * one pass, rather than one pass for each node. */ + num_left = 0; for (k = 0; k < num_isolated; k++) { u = isolated[k]; if (!tsls_hmm_in_tree(self, u)) { - tsls_hmm_delete_likelihood(self, u); L[u] = NOT_IN_TREE; + num_left++; } } + if (num_left > 0) { + tsls_hmm_remove_left_nodes(self); + } } /* Reset the HMM and build the tree containing the specified position, From db2415e93b9e7c3f785a9a4f6feba50eb8b9b3b5 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 13:11:52 +0100 Subject: [PATCH 17/25] Add paired wall-time comparison and document the refactor's results benchmarking/paired.py times each query three times with each engine, alternating, and keeps the best (as in the README's earlier paired measurements), so drift in machine speed doesn't bias one engine. The README now records the instruction counts on the original set after the refactor and optimisations (tsls 0.52-0.97 of tsinfer-hmm, including a whole sample), and the paired wall times, which go the other way: tsls takes 1.14-1.43 times as long. Diagnosis (in the README): building a tsls.Hmm re-encodes every site's alleles (4-8 ms per query against 1.5 ms), and find_path misses the cache more, despite fewer instructions: tsls has 25% more L1 read misses, 2.2x the L1 write misses and 35% more last-level read misses on the profiling sample (callgrind --cache-sim). The last-level misses are mostly tskit's tree position scans reading edges through the index arrays, and the tsk_mutation_t reads; the write misses come from the per-edge marks added by 4c56547 and 0cd8ab3. Also marks the earlier 'where the time goes' section as describing the engine before the refactor. --- benchmarking/README.md | 73 +++++++++++++++++++++++++++++ benchmarking/paired.py | 101 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 174 insertions(+) create mode 100644 benchmarking/paired.py diff --git a/benchmarking/README.md b/benchmarking/README.md index 9920135..b7f49f9 100644 --- a/benchmarking/README.md +++ b/benchmarking/README.md @@ -377,6 +377,11 @@ for the whole group divided by the number of haplotypes. ### tsls against tsinfer-hmm: where the time goes +This section describes the engine before the `tsls_hmm_t` refactor (when +it was `ancestor_matcher_t` in `matcher.c`); see +[Instruction counts after the refactor](#instruction-counts-after-the-tsls_hmm_t-refactor) +for the current state. + 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 @@ -445,3 +450,71 @@ over total tsinfer-hmm time (100 group 1485 ancestors, 6 samples): 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. + +### Instruction counts after the `tsls_hmm_t` refactor + +The engine is now `tsls_hmm_t` (`tsls/lib/hmm.c`). It works directly on a +`tsk_treeseq_t`, moves between trees with tskit's `tsk_tree_position_t`, +and finds the first tree of a window by seeking from a null position +instead of applying every edge before it. The git history from `7277b87` +records each optimisation step and the dead ends, with instruction counts. + +`callgrind.py measure` counts instructions in `find_path` only, with setup +run uninstrumented, on a profiling set sized to run in a few minutes +(tsinfer-hmm's counts are cached, so only tsls is remeasured): + +``` +uv run python callgrind.py measure -j 4 +``` + +The sample in that set is matched against its reference cut at site +100,000, because matching only the first sites of the full reference makes +tsinfer-hmm's traceback walk in from the end of the sequence, which a whole +sample doesn't. On the original set (20 ancestors per group, 1 whole +sample), before the refactor and after the optimisations: + +| workload | tsinfer-hmm | tsls before | tsls after | ratio after | +|---|---|---|---|---| +| ancestors, group 84 | 0.296 G | 0.608 G | 0.154 G | 0.52 | +| ancestors, group 1000 | 0.621 G | 0.975 G | 0.320 G | 0.52 | +| ancestors, group 1483 | 2.388 G | 2.845 G | 1.813 G | 0.76 | +| ancestors, group 1484 | 9.911 G | 11.137 G | 9.046 G | 0.91 | +| ancestors, group 1485 | 6.653 G | 7.445 G | 5.679 G | 0.85 | +| 1 sample | 9.922 G | 10.955 G | 9.633 G | 0.97 | + +tsls finds tsinfer's path for all 1,914 queries in the six workloads, and +its traceback sizes are unchanged. + +### Wall time after the refactor + +Fewer instructions did not mean less time. `paired.py` times each query +three times with each engine, alternating, and keeps the best time for +each (on an idle machine; the ratio is total tsls time over total +tsinfer-hmm time, including building the matcher for each query, as +`engines.py` does): + +| workload | queries | tsinfer-hmm (s) | tsls (s) | ratio | +|---|---|---|---|---| +| ancestors, group 84 | 50 | 0.21 | 0.30 | 1.43 | +| ancestors, group 1000 | 50 | 0.46 | 0.55 | 1.19 | +| ancestors, group 1483 | 100 | 4.76 | 5.67 | 1.19 | +| ancestors, group 1484 | 50 | 5.23 | 6.42 | 1.23 | +| ancestors, group 1485 | 100 | 12.98 | 16.20 | 1.25 | +| samples | 6 | 17.69 | 20.22 | 1.14 | + +Two things account for this: + +- Building a `tsls.Hmm` takes 4-8 ms against 1.5 ms for tsinfer-hmm's + `AncestorMatcher`, because it encodes the alleles of every site from the + tables each time. This dominates for early ancestors. +- `find_path` itself is slower (1.10 on group 1485, 1.18 on samples, timed + separately) despite executing fewer instructions, because it misses the + cache more. Callgrind's cache simulation on the profiling sample: tsls + has 2% fewer instructions but 25% more L1 read misses, 2.2 times the L1 + write misses and 35% more last-level read misses than tsinfer-hmm. The + last-level misses are mostly in tskit's tree position scans, which read + edge coordinates through the index arrays (`left[order[j]]`) where + tsinfer-hmm scans a contiguous array of sorted edges, and in reading + mutations from `tsk_mutation_t` structs. The extra L1 write misses come + from two of the optimisations, which write per-edge marks into + `allelic_state` and `likelihood_cache`. diff --git a/benchmarking/paired.py b/benchmarking/paired.py new file mode 100644 index 0000000..a00354e --- /dev/null +++ b/benchmarking/paired.py @@ -0,0 +1,101 @@ +""" +Compare the wall time of tsls and tsinfer-hmm on the tsinfer workloads with +paired runs. + +run.py runs all the queries with one engine and then all with the next, so +drift in machine speed shows up as a difference between the engines. Here +each query is matched several times with each engine, alternating between +them, and the best time for each engine is kept. The ratio for a workload is +the total tsls time over the total tsinfer-hmm time. Run this on an +otherwise idle machine. +""" + +import pathlib +import time + +import click +import numpy as np + +import callgrind +import queries +import run as run_module + +ENGINES = ["tsinfer", "tsls"] + +DEFAULT_WORKLOADS = [ + ("ancestors_g84", 50), + ("ancestors_g1000", 50), + ("ancestors_g1483", 100), + ("ancestors_g1484", 50), + ("ancestors_g1485", 100), + ("samples", 6), +] + + +def time_workload(workload_dir, num_queries, repeats): + """ + Return the (best_times, same_paths) tuple for the first num_queries + queries, where best_times[engine] is an array of the best time for each + query, and same_paths is True if the engines found the same paths. + """ + model = queries.load_model(workload_dir) + ts = run_module.load_ts(workload_dir / "reference.trees") + H, _, df = queries.load_queries(workload_dir) + engines = {name: callgrind.make_engine(name, ts, model) for name in ENGINES} + query_ids = list(df.query_id[:num_queries]) + best = {name: np.full(len(query_ids), np.inf) for name in ENGINES} + same_paths = True + for j, query_id in enumerate(query_ids): + start, end = queries.query_window(H, query_id) + h = H[query_id] + parents = {} + for rep in range(repeats): + # Alternate which engine goes first + order = ENGINES if rep % 2 == 0 else ENGINES[::-1] + for name in order: + before = time.perf_counter() + _, _, parent, _, _ = engines[name].match(h, start, end) + elapsed = time.perf_counter() - before + best[name][j] = min(best[name][j], elapsed) + parents[name] = list(parent) + same_paths = same_paths and parents["tsinfer"] == parents["tsls"] + return best, same_paths + + +@click.command() +@click.option( + "--data-dir", + type=click.Path(exists=True, file_okay=False), + default="data/tgp_chr20p", +) +@click.option("-w", "--workload", "workload_names", multiple=True) +@click.option("-r", "--repeats", type=int, default=3) +def main(data_dir, workload_names, repeats): + """ + Print a markdown table of paired wall times for each workload. + """ + workloads = DEFAULT_WORKLOADS + if len(workload_names) > 0: + workloads = [w for w in DEFAULT_WORKLOADS if w[0] in workload_names] + click.echo( + "| workload | queries | tsinfer total (s) | tsls total (s) | ratio " + "| median per-query ratio |" + ) + click.echo("|---|---|---|---|---|---|") + for workload, num_queries in workloads: + best, same_paths = time_workload( + pathlib.Path(data_dir) / workload, num_queries, repeats + ) + total = {name: float(np.sum(best[name])) for name in ENGINES} + median = float(np.median(best["tsls"] / best["tsinfer"])) + click.echo( + f"| {workload} | {num_queries} | {total['tsinfer']:.2f} | " + f"{total['tsls']:.2f} | {total['tsls'] / total['tsinfer']:.3f} " + f"| {median:.3f} |", + ) + if not same_paths: + click.echo(f"WARNING: engines found different paths for {workload}") + + +if __name__ == "__main__": + main() From 92d471b977cf73022d2725ee5b9be0fcb8a89245 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 13:41:18 +0100 Subject: [PATCH 18/25] Share an index of the tree sequence between HMMs Paired timing showed that building a tsls.Hmm took 4-8 ms per query against tsinfer-hmm's 1.5 ms, because every HMM encoded the alleles of every site from the tables, which dominated for short ancestors. The read-only data the HMM derives from a tree sequence is now a tsls_index_t, built once (tsls_index_init) and borrowed by any number of HMMs (tsls_hmm_alloc takes the index). It holds the allele encoding and compact per-site mutation arrays (node and encoded derived state, with site offsets), which replace reading the 80-byte tsk_mutation_t structs. In Python, tsls.TreeSequence(ts, alleles=None) owns the tsk_treeseq_t and its index, and Hmm(ts, ..., alleles=...) builds one only when given a tskit tree sequence (alleles with a tsls.TreeSequence is an error). The benchmark engine builds its tsls.TreeSequence once per workload. Cache simulation (callgrind --cache-sim, find_path, profiling sample), before -> after: last-level read misses 1.53M -> 1.31M (tsinfer-hmm 1.14M); L1 read misses 78.9M -> 78.8M (63.0M); L1 write misses unchanged at 16.5-16.9M (7.7M). Instructions 1.525G -> 1.526G. Paired wall time (paired.py --quick, tsls/tsinfer-hmm): ancestors g84 0.735 (1.43 before, on 50 queries), g1485 1.143, samples 1.35 on 2 samples. Paired ratios on samples vary by about +-3% between runs of the same build (1.23 and 1.17 on 6 samples), so cache-miss counts are used to judge the next steps. Instruction counts are unchanged: | workload | queries | tsinfer | tsls | ratio | |---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.155 G | 0.523 | | ancestors_g1000 | 20 | 0.621 G | 0.321 G | 0.517 | | ancestors_g1483 | 10 | 1.355 G | 1.068 G | 0.788 | | ancestors_g1484 | 3 | 0.646 G | 0.568 G | 0.878 | | ancestors_g1485 | 5 | 1.131 G | 0.919 G | 0.813 | | samples | 1 (cut at 100k sites) | 1.561 G | 1.526 G | 0.977 | --- benchmarking/engines.py | 3 +- benchmarking/paired.py | 12 +- experiments/compression/capture.py | 2 +- experiments/thresholding/thresholding.py | 6 +- tsls/_tslsmodule.c | 95 ++++--- tsls/lib/hmm.c | 300 +++++++++++++---------- tsls/lib/tests/tests.c | 51 ++-- tsls/lib/tsls.h | 54 +++- tsls/tests/test_matching.py | 8 +- tsls/tsls/matching.py | 41 ++-- verification/verification.py | 3 +- 11 files changed, 355 insertions(+), 220 deletions(-) diff --git a/benchmarking/engines.py b/benchmarking/engines.py index 48636b4..bc56739 100644 --- a/benchmarking/engines.py +++ b/benchmarking/engines.py @@ -253,14 +253,13 @@ class TslsEngine(AncestorMatcherEngine): def make_reference(self, ts, alleles): tables = ts.dump_tables() tables.edges.keep_rows(tables.edges.parent != 0) - return tsls.TreeSequence(tables.tree_sequence()) + return tsls.TreeSequence(tables.tree_sequence(), alleles=alleles) def make_matcher(self): return tsls.Hmm( self.reference, self.recombination, self.mismatch, - alleles=self.alleles, likelihood_threshold=self.likelihood_threshold, weight_by_n=self.weight_by_n, ) diff --git a/benchmarking/paired.py b/benchmarking/paired.py index a00354e..3f4164e 100644 --- a/benchmarking/paired.py +++ b/benchmarking/paired.py @@ -31,6 +31,13 @@ ("samples", 6), ] +# A smaller set for checking changes, which takes about two minutes +QUICK_WORKLOADS = [ + ("ancestors_g84", 50), + ("ancestors_g1485", 30), + ("samples", 2), +] + def time_workload(workload_dir, num_queries, repeats): """ @@ -70,11 +77,12 @@ def time_workload(workload_dir, num_queries, repeats): ) @click.option("-w", "--workload", "workload_names", multiple=True) @click.option("-r", "--repeats", type=int, default=3) -def main(data_dir, workload_names, repeats): +@click.option("--quick", is_flag=True, help="Use the smaller set of workloads") +def main(data_dir, workload_names, repeats, quick): """ Print a markdown table of paired wall times for each workload. """ - workloads = DEFAULT_WORKLOADS + workloads = QUICK_WORKLOADS if quick else DEFAULT_WORKLOADS if len(workload_names) > 0: workloads = [w for w in DEFAULT_WORKLOADS if w[0] in workload_names] click.echo( diff --git a/experiments/compression/capture.py b/experiments/compression/capture.py index d809ab6..6089084 100644 --- a/experiments/compression/capture.py +++ b/experiments/compression/capture.py @@ -61,7 +61,7 @@ def sc2ts_queries(): def capture(name, ts, query_list, mu, rho, weight_by_n, rows): ts_path = RESULTS / f"{name}.trees" ts.dump(ts_path) - mi = tsls.TreeSequence(ts) + mi = tsls.TreeSequence(ts, alleles=v.ALLELES) for query_id, label, h in query_list: am = v.make_matcher( ts, mu, rho, weight_by_n=weight_by_n, count_distinct=True, tree_sequence=mi diff --git a/experiments/thresholding/thresholding.py b/experiments/thresholding/thresholding.py index 8f91c2d..86cc496 100644 --- a/experiments/thresholding/thresholding.py +++ b/experiments/thresholding/thresholding.py @@ -206,7 +206,7 @@ def run_threading_task(num_samples, seed, quick): ts, h, true_parent, near = threading_case(num_samples, seed) m = ts.num_sites mu, rho = v.hmm_params(v.THREADING_DEFAULTS["k"]) - mi = tsls.TreeSequence(ts) + mi = tsls.TreeSequence(ts, alleles=v.ALLELES) true_switch = np.where(true_parent[1:] != true_parent[:-1])[0] + 1 rows = [] exact_posterior = None @@ -291,7 +291,7 @@ def run_calibration_task(setting, seed, quick): m = ts.num_sites n = ts.num_nodes rho = switches / m - mi = tsls.TreeSequence(ts) + mi = tsls.TreeSequence(ts, alleles=v.ALLELES) queries = [] for _ in range(3 if quick else 5): path = v.sample_prior_path(present, continuous, np.full(m, rho), n, rng) @@ -363,7 +363,7 @@ def load_sc2ts(size): tables = ts.dump_tables() tables.edges.keep_rows(tables.edges.parent != 0) ts = tables.tree_sequence() - mi = tsls.TreeSequence(ts) + mi = tsls.TreeSequence(ts, alleles=v.ALLELES) query_dir = SC2TS_DATA / "queries_scaling" H = np.load(query_dir / "queries.npz")["H"] df = pd.read_csv(query_dir / "queries.csv") diff --git a/tsls/_tslsmodule.c b/tsls/_tslsmodule.c index 37b6b49..617e84f 100644 --- a/tsls/_tslsmodule.c +++ b/tsls/_tslsmodule.c @@ -24,6 +24,7 @@ static PyObject *TslsMatchImpossible; typedef struct { PyObject_HEAD tsk_treeseq_t *tree_sequence; + tsls_index_t *index; } TreeSequence; typedef struct { @@ -52,13 +53,14 @@ handle_library_error(int err) */ /* A minimal tree sequence type, built from a LightweightTableCollection, - * which can be shared between HMMs (including in different threads). */ + * with the index the HMM needs (for the specified alleles), which can be + * shared between HMMs (including in different threads). */ static int TreeSequence_check_state(TreeSequence *self) { int ret = 0; - if (self->tree_sequence == NULL) { + if (self->tree_sequence == NULL || self->index == NULL) { PyErr_SetString(PyExc_SystemError, "TreeSequence not initialised"); ret = -1; } @@ -68,6 +70,11 @@ TreeSequence_check_state(TreeSequence *self) static void TreeSequence_dealloc(TreeSequence *self) { + if (self->index != NULL) { + tsls_index_free(self->index); + PyMem_Free(self->index); + self->index = NULL; + } if (self->tree_sequence != NULL) { tsk_treeseq_free(self->tree_sequence); PyMem_Free(self->tree_sequence); @@ -82,13 +89,39 @@ TreeSequence_init(TreeSequence *self, PyObject *args, PyObject *kwds) int ret = -1; int err; LightweightTableCollection *tables = NULL; - static char *kwlist[] = { "tables", NULL }; + PyObject *alleles = Py_None; + PyObject *alleles_seq = NULL; + const char **alleles_data = NULL; + tsk_size_t num_alleles = 0; + Py_ssize_t j; + static char *kwlist[] = { "tables", "alleles", NULL }; self->tree_sequence = NULL; - if (!PyArg_ParseTupleAndKeywords( - args, kwds, "O!", kwlist, &LightweightTableCollectionType, &tables)) { + self->index = NULL; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!|O", kwlist, + &LightweightTableCollectionType, &tables, &alleles)) { goto out; } + if (alleles != Py_None) { + alleles_seq = PySequence_Fast(alleles, "alleles must be a sequence"); + if (alleles_seq == NULL) { + goto out; + } + num_alleles = (tsk_size_t) PySequence_Fast_GET_SIZE(alleles_seq); + alleles_data = PyMem_Calloc(num_alleles + 1, sizeof(*alleles_data)); + if (alleles_data == NULL) { + PyErr_NoMemory(); + goto out; + } + for (j = 0; j < (Py_ssize_t) num_alleles; j++) { + /* The UTF-8 buffer lives as long as the string */ + alleles_data[j] + = PyUnicode_AsUTF8(PySequence_Fast_GET_ITEM(alleles_seq, j)); + if (alleles_data[j] == NULL) { + goto out; + } + } + } if (LightweightTableCollection_check_state(tables) != 0) { goto out; } @@ -103,8 +136,20 @@ TreeSequence_init(TreeSequence *self, PyObject *args, PyObject *kwds) handle_tskit_error(err); goto out; } + self->index = PyMem_Calloc(1, sizeof(*self->index)); + if (self->index == NULL) { + PyErr_NoMemory(); + goto out; + } + err = tsls_index_init(self->index, self->tree_sequence, alleles_data, num_alleles); + if (err != 0) { + handle_library_error(err); + goto out; + } ret = 0; out: + PyMem_Free(alleles_data); + Py_XDECREF(alleles_seq); return ret; } @@ -187,27 +232,22 @@ Hmm_init(Hmm *self, PyObject *args, PyObject *kwds) double pruning_parameter = 0; int count_distinct = 0; static char *kwlist[] = { "tree_sequence", "recombination", "mismatch", - "alleles", "extended_checks", "likelihood_threshold", "weight_by_n", - "pruning_mode", "pruning_parameter", "count_distinct", NULL }; + "extended_checks", "likelihood_threshold", "weight_by_n", "pruning_mode", + "pruning_parameter", "count_distinct", NULL }; TreeSequence *tree_sequence = NULL; PyObject *recombination = NULL; PyObject *mismatch = NULL; - PyObject *alleles = Py_None; - PyObject *alleles_seq = NULL; - const char **alleles_data = NULL; - tsk_size_t num_alleles = 0; PyArrayObject *recombination_array = NULL; PyArrayObject *mismatch_array = NULL; npy_intp *shape; double likelihood_threshold = DBL_MIN; int flags = 0; - Py_ssize_t j; size_t num_sites; self->hmm = NULL; self->tree_sequence = NULL; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!OO|Oidiidi", kwlist, - &TreeSequenceType, &tree_sequence, &recombination, &mismatch, &alleles, + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!OO|idiidi", kwlist, + &TreeSequenceType, &tree_sequence, &recombination, &mismatch, &extended_checks, &likelihood_threshold, &weight_by_n, &pruning_mode, &pruning_parameter, &count_distinct)) { goto out; @@ -219,27 +259,6 @@ Hmm_init(Hmm *self, PyObject *args, PyObject *kwds) } num_sites = tsk_treeseq_get_num_sites(tree_sequence->tree_sequence); - if (alleles != Py_None) { - alleles_seq = PySequence_Fast(alleles, "alleles must be a sequence"); - if (alleles_seq == NULL) { - goto out; - } - num_alleles = (tsk_size_t) PySequence_Fast_GET_SIZE(alleles_seq); - alleles_data = PyMem_Calloc(num_alleles + 1, sizeof(*alleles_data)); - if (alleles_data == NULL) { - PyErr_NoMemory(); - goto out; - } - for (j = 0; j < (Py_ssize_t) num_alleles; j++) { - /* The UTF-8 buffer lives as long as the string */ - alleles_data[j] - = PyUnicode_AsUTF8(PySequence_Fast_GET_ITEM(alleles_seq, j)); - if (alleles_data[j] == NULL) { - goto out; - } - } - } - recombination_array = (PyArrayObject *) PyArray_FromAny(recombination, PyArray_DescrFromType(NPY_FLOAT64), 1, 1, NPY_ARRAY_IN_ARRAY, NULL); if (recombination_array == NULL) { @@ -276,8 +295,8 @@ Hmm_init(Hmm *self, PyObject *args, PyObject *kwds) if (count_distinct) { flags |= TSLS_COUNT_DISTINCT; } - err = tsls_hmm_alloc(self->hmm, self->tree_sequence->tree_sequence, alleles_data, - num_alleles, PyArray_DATA(recombination_array), PyArray_DATA(mismatch_array), + err = tsls_hmm_alloc(self->hmm, self->tree_sequence->index, + PyArray_DATA(recombination_array), PyArray_DATA(mismatch_array), likelihood_threshold, flags); if (err != 0) { handle_library_error(err); @@ -290,8 +309,6 @@ Hmm_init(Hmm *self, PyObject *args, PyObject *kwds) } ret = 0; out: - PyMem_Free(alleles_data); - Py_XDECREF(alleles_seq); Py_XDECREF(recombination_array); Py_XDECREF(mismatch_array); return ret; diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index 5791567..5f86049 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -248,87 +248,10 @@ compressed_matrix_store_site(compressed_matrix_t *self, tsk_id_t site, } /*=================================================================== - * HMM + * Index *=================================================================== */ -static inline bool -tsls_hmm_in_tree(const tsls_hmm_t *self, tsk_id_t u) -{ - return self->parent[u] != NULL_NODE || self->left_child[u] != NULL_NODE; -} - -static void -tsls_hmm_check_state(const tsls_hmm_t *self) -{ - int num_likelihoods = 0; - int j; - tsk_id_t u; - - for (j = 0; j < self->num_likelihood_nodes; j++) { - u = self->likelihood_nodes[j]; - assert(self->likelihood[u] >= 0); - } - for (u = 0; u < (tsk_id_t) self->num_nodes; u++) { - if (self->likelihood[u] >= 0) { - num_likelihoods++; - } - if (tsls_hmm_in_tree(self, u)) { - assert(self->likelihood[u] != NOT_IN_TREE); - if (self->parent[u] == NULL_NODE) { - /* Roots always have a likelihood */ - assert(self->likelihood[u] >= 0); - } - } else { - assert(self->likelihood[u] == NOT_IN_TREE); - } - assert(self->allelic_state[u] - == (self->parent[u] == NULL_NODE ? ROOT_ALLELE : TSK_NULL)); - assert(self->likelihood_cache[u] == CACHE_UNSET); - } - assert(num_likelihoods == self->num_likelihood_nodes); - (void) num_likelihoods; -} - -int -tsls_hmm_print_state(tsls_hmm_t *self, FILE *out) -{ - int j, k; - tsk_id_t u; - - fprintf(out, "HMM state\n"); - fprintf(out, "site\tposition\trecomb_rate\tmut_rate\tnum_alleles\tancestral\n"); - for (j = 0; j < (int) self->num_sites; j++) { - fprintf(out, "%d\t%f\t%f\t%f\t%d\t%d\n", j, self->sites_position[j], - self->recombination_rate[j], self->mismatch_rate[j], - (int) self->num_alleles[j], (int) self->ancestral_state[j]); - } - fprintf(out, "tree = \n"); - fprintf(out, "id\tparent\tlchild\trchild\tlsib\trsib\tlikelihood\n"); - for (j = 0; j < (int) self->num_nodes; j++) { - fprintf(out, "%d\t%d\t%d\t%d\t%d\t%d\t%f\n", (int) j, self->parent[j], - self->left_child[j], self->right_child[j], self->left_sib[j], - self->right_sib[j], self->likelihood[j]); - } - fprintf(out, "likelihood nodes\n"); - for (j = 0; j < self->num_likelihood_nodes; j++) { - u = self->likelihood_nodes[j]; - fprintf(out, "\t%d -> %f\n", u, self->likelihood[u]); - } - fprintf(out, "traceback\n"); - for (j = 0; j < (int) self->num_sites; j++) { - fprintf(out, "\t%d:%d (%d)\t", (int) j, self->max_likelihood_node[j], - self->traceback[j].size); - for (k = 0; k < self->traceback[j].size; k++) { - fprintf(out, "(%d, %d)", self->traceback[j].node[k], - self->traceback[j].recombination_required[k]); - } - fprintf(out, "\n"); - } - tsi_blkalloc_print_state(&self->traceback_allocator, out); - return 0; -} - /* Returns the index of the specified state in the list of alleles, or -1 * if it is not in the list. */ static int @@ -375,8 +298,8 @@ encode_allele(const char *state, tsk_size_t length, const char **alleles, * two alleles at each site. Otherwise, each * state is its index in the specified list of alleles. */ static int WARN_UNUSED -tsls_hmm_encode_alleles( - tsls_hmm_t *self, const char **alleles, tsk_size_t num_global_alleles) +tsls_index_encode_alleles( + tsls_index_t *self, const char **alleles, tsk_size_t num_global_alleles) { int ret = 0; const tsk_treeseq_t *ts = self->tree_sequence; @@ -387,6 +310,7 @@ tsls_hmm_encode_alleles( const char **list = global ? alleles : site_alleles; const tsk_mutation_t *mutation; tsk_size_t j, k, n, offset; + tsk_size_t mutation_offset = 0; int code; if (global) { @@ -399,6 +323,10 @@ tsls_hmm_encode_alleles( } } for (j = 0; j < self->num_sites; j++) { + /* Mutations are sorted by site, so the mutations of each site are + * contiguous in the mutation table */ + self->site_mutations_offset[j] = mutation_offset; + mutation_offset += ts->site_mutations_length[j]; n = global ? num_global_alleles : 0; offset = sites->ancestral_state_offset[j]; code = encode_allele(sites->ancestral_state + offset, @@ -417,38 +345,168 @@ tsls_hmm_encode_alleles( ret = code; goto out; } - self->derived_state[mutation->id] = (allele_t) code; + assert(mutation->id == (tsk_id_t) (self->site_mutations_offset[j] + k)); + self->mutation_node[mutation->id] = mutation->node; + self->mutation_derived_state[mutation->id] = (allele_t) code; } /* A site always has at least two alleles, so that a haplotype can * differ from a site with no mutations (such as a tsinfer * ancestor's focal sites, which are not yet in the tree sequence). */ self->num_alleles[j] = TSK_MAX(n, 2); } + self->site_mutations_offset[self->num_sites] = mutation_offset; out: return ret; } int -tsls_hmm_alloc(tsls_hmm_t *self, const tsk_treeseq_t *tree_sequence, - const char **alleles, tsk_size_t num_alleles, const double *recombination_rate, - const double *mismatch_rate, double likelihood_threshold, int flags) +tsls_index_init(tsls_index_t *self, const tsk_treeseq_t *tree_sequence, + const char **alleles, tsk_size_t num_alleles) { int ret = 0; - const size_t N = tsk_treeseq_get_num_nodes(tree_sequence); const size_t num_mutations = tsk_treeseq_get_num_mutations(tree_sequence); - memset(self, 0, sizeof(tsls_hmm_t)); - self->flags = flags; - self->likelihood_threshold = likelihood_threshold; + memset(self, 0, sizeof(*self)); self->tree_sequence = tree_sequence; + self->num_nodes = tsk_treeseq_get_num_nodes(tree_sequence); self->num_sites = tsk_treeseq_get_num_sites(tree_sequence); - self->num_nodes = N; + self->num_edges = tsk_treeseq_get_num_edges(tree_sequence); self->sites_position = tree_sequence->tables->sites.position; self->sequence_length = tsk_treeseq_get_sequence_length(tree_sequence); self->num_alleles = malloc((self->num_sites + 1) * sizeof(*self->num_alleles)); self->ancestral_state = malloc((self->num_sites + 1) * sizeof(*self->ancestral_state)); - self->derived_state = malloc((num_mutations + 1) * sizeof(*self->derived_state)); + self->site_mutations_offset + = malloc((self->num_sites + 1) * sizeof(*self->site_mutations_offset)); + self->mutation_node = malloc((num_mutations + 1) * sizeof(*self->mutation_node)); + self->mutation_derived_state + = malloc((num_mutations + 1) * sizeof(*self->mutation_derived_state)); + if (self->num_alleles == NULL || self->ancestral_state == NULL + || self->site_mutations_offset == NULL || self->mutation_node == NULL + || self->mutation_derived_state == NULL) { + ret = TSLS_ERR_NO_MEMORY; + goto out; + } + ret = tsls_index_encode_alleles(self, alleles, num_alleles); + if (ret != 0) { + goto out; + } +out: + return ret; +} + +int +tsls_index_free(tsls_index_t *self) +{ + tsi_safe_free(self->num_alleles); + tsi_safe_free(self->ancestral_state); + tsi_safe_free(self->site_mutations_offset); + tsi_safe_free(self->mutation_node); + tsi_safe_free(self->mutation_derived_state); + tsi_safe_free(self->insertion_edges); + tsi_safe_free(self->removal_edges); + return 0; +} + +/*=================================================================== + * HMM + *=================================================================== + */ + +static inline bool +tsls_hmm_in_tree(const tsls_hmm_t *self, tsk_id_t u) +{ + return self->parent[u] != NULL_NODE || self->left_child[u] != NULL_NODE; +} + +static void +tsls_hmm_check_state(const tsls_hmm_t *self) +{ + int num_likelihoods = 0; + int j; + tsk_id_t u; + + for (j = 0; j < self->num_likelihood_nodes; j++) { + u = self->likelihood_nodes[j]; + assert(self->likelihood[u] >= 0); + } + for (u = 0; u < (tsk_id_t) self->num_nodes; u++) { + if (self->likelihood[u] >= 0) { + num_likelihoods++; + } + if (tsls_hmm_in_tree(self, u)) { + assert(self->likelihood[u] != NOT_IN_TREE); + if (self->parent[u] == NULL_NODE) { + /* Roots always have a likelihood */ + assert(self->likelihood[u] >= 0); + } + } else { + assert(self->likelihood[u] == NOT_IN_TREE); + } + assert(self->allelic_state[u] + == (self->parent[u] == NULL_NODE ? ROOT_ALLELE : TSK_NULL)); + assert(self->likelihood_cache[u] == CACHE_UNSET); + } + assert(num_likelihoods == self->num_likelihood_nodes); + (void) num_likelihoods; +} + +int +tsls_hmm_print_state(tsls_hmm_t *self, FILE *out) +{ + int j, k; + tsk_id_t u; + + fprintf(out, "HMM state\n"); + fprintf(out, "site\tposition\trecomb_rate\tmut_rate\tnum_alleles\tancestral\n"); + for (j = 0; j < (int) self->num_sites; j++) { + fprintf(out, "%d\t%f\t%f\t%f\t%d\t%d\n", j, self->sites_position[j], + self->recombination_rate[j], self->mismatch_rate[j], + (int) self->index->num_alleles[j], (int) self->index->ancestral_state[j]); + } + fprintf(out, "tree = \n"); + fprintf(out, "id\tparent\tlchild\trchild\tlsib\trsib\tlikelihood\n"); + for (j = 0; j < (int) self->num_nodes; j++) { + fprintf(out, "%d\t%d\t%d\t%d\t%d\t%d\t%f\n", (int) j, self->parent[j], + self->left_child[j], self->right_child[j], self->left_sib[j], + self->right_sib[j], self->likelihood[j]); + } + fprintf(out, "likelihood nodes\n"); + for (j = 0; j < self->num_likelihood_nodes; j++) { + u = self->likelihood_nodes[j]; + fprintf(out, "\t%d -> %f\n", u, self->likelihood[u]); + } + fprintf(out, "traceback\n"); + for (j = 0; j < (int) self->num_sites; j++) { + fprintf(out, "\t%d:%d (%d)\t", (int) j, self->max_likelihood_node[j], + self->traceback[j].size); + for (k = 0; k < self->traceback[j].size; k++) { + fprintf(out, "(%d, %d)", self->traceback[j].node[k], + self->traceback[j].recombination_required[k]); + } + fprintf(out, "\n"); + } + tsi_blkalloc_print_state(&self->traceback_allocator, out); + return 0; +} + +int +tsls_hmm_alloc(tsls_hmm_t *self, const tsls_index_t *index, + const double *recombination_rate, const double *mismatch_rate, + double likelihood_threshold, int flags) +{ + int ret = 0; + const size_t N = index->num_nodes; + + memset(self, 0, sizeof(tsls_hmm_t)); + self->flags = flags; + self->likelihood_threshold = likelihood_threshold; + self->index = index; + self->tree_sequence = index->tree_sequence; + self->num_sites = index->num_sites; + self->num_nodes = N; + self->sites_position = index->sites_position; + self->sequence_length = index->sequence_length; self->recombination_rate = malloc((self->num_sites + 1) * sizeof(*self->recombination_rate)); self->mismatch_rate = malloc((self->num_sites + 1) * sizeof(*self->mismatch_rate)); @@ -473,8 +531,7 @@ tsls_hmm_alloc(tsls_hmm_t *self, const tsk_treeseq_t *tree_sequence, self->likelihood_nodes_tmp = malloc((N + 1) * sizeof(*self->likelihood_nodes_tmp)); self->values_tmp = malloc((N + 1) * sizeof(*self->values_tmp)); - if (self->num_alleles == NULL || self->ancestral_state == NULL - || self->derived_state == NULL || self->recombination_rate == NULL + if (self->recombination_rate == NULL || self->mismatch_rate == NULL || self->traceback == NULL || self->max_likelihood_node == NULL || self->parent == NULL || self->left_child == NULL || self->right_child == NULL @@ -487,10 +544,6 @@ tsls_hmm_alloc(tsls_hmm_t *self, const tsk_treeseq_t *tree_sequence, ret = TSLS_ERR_NO_MEMORY; goto out; } - ret = tsls_hmm_encode_alleles(self, alleles, num_alleles); - if (ret != 0) { - goto out; - } /* Alloc in blocks of at least 64MiB, and large enough to hold the * traceback for a site in which every node has a likelihood. */ self->traceback_block_size @@ -516,9 +569,6 @@ tsls_hmm_alloc(tsls_hmm_t *self, const tsk_treeseq_t *tree_sequence, int tsls_hmm_free(tsls_hmm_t *self) { - tsi_safe_free(self->num_alleles); - tsi_safe_free(self->ancestral_state); - tsi_safe_free(self->derived_state); tsi_safe_free(self->recombination_rate); tsi_safe_free(self->mismatch_rate); tsi_safe_free(self->parent); @@ -893,12 +943,14 @@ tsls_hmm_init_sweep( static inline void tsls_hmm_set_allelic_state(tsls_hmm_t *self, const tsk_id_t site) { - const tsk_mutation_t *restrict mutations = self->tree_sequence->site_mutations[site]; - const tsk_size_t num_mutations = self->tree_sequence->site_mutations_length[site]; + const tsls_index_t *index = self->index; + const tsk_id_t *restrict node = index->mutation_node; + const allele_t *restrict derived_state = index->mutation_derived_state; + const tsk_size_t stop = index->site_mutations_offset[site + 1]; tsk_size_t j; - for (j = 0; j < num_mutations; j++) { - self->allelic_state[mutations[j].node] = self->derived_state[mutations[j].id]; + for (j = index->site_mutations_offset[site]; j < stop; j++) { + self->allelic_state[node[j]] = derived_state[j]; } } @@ -907,13 +959,14 @@ tsls_hmm_set_allelic_state(tsls_hmm_t *self, const tsk_id_t site) static inline void tsls_hmm_unset_allelic_state(tsls_hmm_t *self, const tsk_id_t site) { - const tsk_mutation_t *restrict mutations = self->tree_sequence->site_mutations[site]; - const tsk_size_t num_mutations = self->tree_sequence->site_mutations_length[site]; + const tsls_index_t *index = self->index; + const tsk_id_t *restrict node = index->mutation_node; + const tsk_size_t stop = index->site_mutations_offset[site + 1]; tsk_size_t j; tsk_id_t u; - for (j = 0; j < num_mutations; j++) { - u = mutations[j].node; + for (j = index->site_mutations_offset[site]; j < stop; j++) { + u = node[j]; self->allelic_state[u] = self->parent[u] == NULL_NODE ? ROOT_ALLELE : TSK_NULL; } } @@ -923,12 +976,13 @@ tsls_hmm_unset_allelic_state(tsls_hmm_t *self, const tsk_id_t site) static inline void tsls_hmm_clear_allelic_state(tsls_hmm_t *self, const tsk_id_t site) { - const tsk_mutation_t *restrict mutations = self->tree_sequence->site_mutations[site]; - const tsk_size_t num_mutations = self->tree_sequence->site_mutations_length[site]; + const tsls_index_t *index = self->index; + const tsk_id_t *restrict node = index->mutation_node; + const tsk_size_t stop = index->site_mutations_offset[site + 1]; tsk_size_t j; - for (j = 0; j < num_mutations; j++) { - self->allelic_state[mutations[j].node] = TSK_NULL; + for (j = index->site_mutations_offset[site]; j < stop; j++) { + self->allelic_state[node[j]] = TSK_NULL; } } @@ -939,13 +993,13 @@ tsls_hmm_insert_mutation_likelihoods(tsls_hmm_t *self, const tsk_id_t site) { const double *restrict L = self->likelihood; const tsk_id_t *restrict parent = self->parent; - const tsk_mutation_t *restrict mutations = self->tree_sequence->site_mutations[site]; - const tsk_size_t num_mutations = self->tree_sequence->site_mutations_length[site]; + const tsls_index_t *index = self->index; + const tsk_size_t stop = index->site_mutations_offset[site + 1]; tsk_size_t j; tsk_id_t u, node; - for (j = 0; j < num_mutations; j++) { - node = mutations[j].node; + for (j = index->site_mutations_offset[site]; j < stop; j++) { + node = index->mutation_node[j]; if (L[node] == NULL_LIKELIHOOD) { u = node; while (L[u] == NULL_LIKELIHOOD) { @@ -1059,7 +1113,7 @@ emission_probability(allele_t node_state, allele_t state, double mu, double num_ static inline int check_haplotype_state(const tsls_hmm_t *self, tsk_id_t site, allele_t state) { - const tsk_size_t num_alleles = self->num_alleles[site]; + const tsk_size_t num_alleles = self->index->num_alleles[site]; if (state < TSLS_MISSING_DATA || (state >= 0 && (tsk_size_t) state >= num_alleles)) { return TSLS_ERR_BAD_HAPLOTYPE_ALLELE; @@ -1200,8 +1254,8 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site, const double rho = self->recombination_rate[site]; const double mu = self->mismatch_rate[site]; const double n = tsls_hmm_get_n(self); - const double num_alleles = (double) self->num_alleles[site]; - const allele_t ancestral_state = self->ancestral_state[site]; + const double num_alleles = (double) self->index->num_alleles[site]; + const allele_t ancestral_state = self->index->ancestral_state[site]; double max_L, p_no_recomb, p_recomb, p_t, p_e; tsk_id_t u, max_L_node; int j; @@ -1286,8 +1340,8 @@ tsls_hmm_process_site_forward(tsls_hmm_t *self, const tsk_id_t site, const double rho = self->recombination_rate[site]; const double mu = self->mismatch_rate[site]; const double n = tsls_hmm_get_n(self); - const double num_alleles = (double) self->num_alleles[site]; - const allele_t ancestral_state = self->ancestral_state[site]; + const double num_alleles = (double) self->index->num_alleles[site]; + const allele_t ancestral_state = self->index->ancestral_state[site]; double total, normalisation_factor, p_e; tsk_id_t u; int j; @@ -1351,8 +1405,8 @@ tsls_hmm_process_site_backward(tsls_hmm_t *self, const tsk_id_t site, double *restrict L = self->likelihood; const double mu = self->mismatch_rate[site]; const double n = tsls_hmm_get_n(self); - const double num_alleles = (double) self->num_alleles[site]; - const allele_t ancestral_state = self->ancestral_state[site]; + const double num_alleles = (double) self->index->num_alleles[site]; + const allele_t ancestral_state = self->index->ancestral_state[site]; double x; tsk_id_t u; int j; @@ -1553,7 +1607,7 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, parent[edge_child[e]] = edge_parent[e]; } } - allelic_state[virtual_root] = self->ancestral_state[site]; + allelic_state[virtual_root] = self->index->ancestral_state[site]; tsls_hmm_set_allelic_state(self, site); u = path_parent[path_length]; while (allelic_state[u] == TSK_NULL) { diff --git a/tsls/lib/tests/tests.c b/tsls/lib/tests/tests.c index fd40a73..df5c043 100644 --- a/tsls/lib/tests/tests.c +++ b/tsls/lib/tests/tests.c @@ -77,7 +77,7 @@ test_allele_encoding(void) { int ret; tsk_treeseq_t ts; - tsls_hmm_t hmm; + tsls_index_t index; tsk_size_t j; tsk_mutation_t *mutation; const char *too_few[] = { "A", "C" }; @@ -87,43 +87,48 @@ test_allele_encoding(void) build_forest_ts(&ts, true); - ret = tsls_hmm_alloc(&hmm, &ts, acgt, 4, rho, mu, DBL_MIN, 0); + ret = tsls_index_init(&index, &ts, acgt, 4); CU_ASSERT_EQUAL_FATAL(ret, 0); for (j = 0; j < tsk_treeseq_get_num_sites(&ts); j++) { - CU_ASSERT_EQUAL(hmm.num_alleles[j], 4); - CU_ASSERT_EQUAL(hmm.ancestral_state[j], + CU_ASSERT_EQUAL(index.num_alleles[j], 4); + CU_ASSERT_EQUAL(index.ancestral_state[j], acgt_code(ts.tables->sites .ancestral_state[ts.tables->sites.ancestral_state_offset[j]])); } for (j = 0; j < tsk_treeseq_get_num_mutations(&ts); j++) { mutation = &ts.site_mutations_mem[j]; - CU_ASSERT_EQUAL( - hmm.derived_state[mutation->id], acgt_code(mutation->derived_state[0])); + CU_ASSERT_EQUAL(index.mutation_node[mutation->id], mutation->node); + CU_ASSERT_EQUAL(index.mutation_derived_state[mutation->id], + acgt_code(mutation->derived_state[0])); } - tsls_hmm_free(&hmm); + for (j = 0; j < tsk_treeseq_get_num_sites(&ts); j++) { + CU_ASSERT_EQUAL(index.site_mutations_offset[j + 1] - index.site_mutations_offset[j], + ts.site_mutations_length[j]); + } + tsls_index_free(&index); /* Per-site encoding, as in tskit */ - ret = tsls_hmm_alloc(&hmm, &ts, NULL, 0, rho, mu, DBL_MIN, 0); + ret = tsls_index_init(&index, &ts, NULL, 0); CU_ASSERT_EQUAL_FATAL(ret, 0); for (j = 0; j < tsk_treeseq_get_num_sites(&ts); j++) { - CU_ASSERT_EQUAL(hmm.num_alleles[j], site_num_alleles[j]); - CU_ASSERT_EQUAL(hmm.ancestral_state[j], 0); + CU_ASSERT_EQUAL(index.num_alleles[j], site_num_alleles[j]); + CU_ASSERT_EQUAL(index.ancestral_state[j], 0); } /* At the last site, the mutation to A is a back mutation */ for (j = 0; j < ts.site_mutations_length[3]; j++) { mutation = &ts.site_mutations[3][j]; - CU_ASSERT_EQUAL( - hmm.derived_state[mutation->id], mutation->derived_state[0] == 'A' ? 0 : 1); + CU_ASSERT_EQUAL(index.mutation_derived_state[mutation->id], + mutation->derived_state[0] == 'A' ? 0 : 1); } - tsls_hmm_free(&hmm); + tsls_index_free(&index); - ret = tsls_hmm_alloc(&hmm, &ts, too_few, 2, rho, mu, DBL_MIN, 0); + ret = tsls_index_init(&index, &ts, too_few, 2); CU_ASSERT_EQUAL(ret, TSLS_ERR_UNKNOWN_ALLELE); - tsls_hmm_free(&hmm); + tsls_index_free(&index); - ret = tsls_hmm_alloc(&hmm, &ts, acgt, TSLS_MAX_ALLELES + 1, rho, mu, DBL_MIN, 0); + ret = tsls_index_init(&index, &ts, acgt, TSLS_MAX_ALLELES + 1); CU_ASSERT_EQUAL(ret, TSLS_ERR_TOO_MANY_ALLELES); - tsls_hmm_free(&hmm); + tsls_index_free(&index); tsk_treeseq_free(&ts); } @@ -133,6 +138,7 @@ test_forest_matching(void) { int ret; tsk_treeseq_t ts; + tsls_index_t index; tsls_hmm_t hmm; compressed_matrix_t fm, bm; /* The haplotype of node 1: C (under 3), C (ancestral), G, A (back mutation) */ @@ -144,7 +150,9 @@ test_forest_matching(void) double total; build_forest_ts(&ts, false); - ret = tsls_hmm_alloc(&hmm, &ts, acgt, 4, rho, mu, DBL_MIN, TSLS_EXTENDED_CHECKS); + ret = tsls_index_init(&index, &ts, acgt, 4); + CU_ASSERT_EQUAL_FATAL(ret, 0); + ret = tsls_hmm_alloc(&hmm, &index, rho, mu, DBL_MIN, TSLS_EXTENDED_CHECKS); CU_ASSERT_EQUAL_FATAL(ret, 0); ret = tsls_hmm_find_path(&hmm, 0, 4, h, match, &path_length, left, right, parent); @@ -214,6 +222,7 @@ test_forest_matching(void) compressed_matrix_free(&bm); tsls_hmm_free(&hmm); + tsls_index_free(&index); tsk_treeseq_free(&ts); } @@ -222,6 +231,7 @@ test_bad_inputs(void) { int ret; tsk_treeseq_t ts; + tsls_index_t index; tsls_hmm_t hmm; compressed_matrix_t cm; double zero_mu[] = { 0, 0, 0, 0 }; @@ -235,7 +245,9 @@ test_bad_inputs(void) size_t path_length; build_forest_ts(&ts, false); - ret = tsls_hmm_alloc(&hmm, &ts, acgt, 4, rho, zero_mu, DBL_MIN, 0); + ret = tsls_index_init(&index, &ts, acgt, 4); + CU_ASSERT_EQUAL_FATAL(ret, 0); + ret = tsls_hmm_alloc(&hmm, &index, rho, zero_mu, DBL_MIN, 0); CU_ASSERT_EQUAL_FATAL(ret, 0); ret = tsls_hmm_find_path( @@ -260,6 +272,7 @@ test_bad_inputs(void) compressed_matrix_free(&cm); tsls_hmm_free(&hmm); + tsls_index_free(&index); tsk_treeseq_free(&ts); } diff --git a/tsls/lib/tsls.h b/tsls/lib/tsls.h index 739b423..38ad3c4 100644 --- a/tsls/lib/tsls.h +++ b/tsls/lib/tsls.h @@ -92,20 +92,54 @@ typedef struct { tsi_blkalloc_t allocator; } compressed_matrix_t; -/* The Li and Stephens HMM on a tree sequence. The tree sequence is borrowed, - * and must not be changed or freed while the HMM is in use. */ +/* An edge in one of tskit's edge orders, with the one coordinate that + * finding the first tree needs: right in insertion order (in which the + * edges before a position all have left <= it), and left in removal + * order. */ +typedef struct { + tsk_id_t parent; + tsk_id_t child; + double coordinate; +} tsls_edge_t; + +/* Read-only data derived from a tree sequence for the HMM, which is built + * once and can be shared between HMMs (including in different threads). + * The tree sequence is borrowed, and must not be changed or freed while the + * index is in use. */ typedef struct { - int flags; const tsk_treeseq_t *tree_sequence; size_t num_nodes; size_t num_sites; + size_t num_edges; const double *sites_position; double sequence_length; /* The integer encoding of the alleles: the number of alleles and the - * ancestral state of each site, and the derived state of each mutation */ + * ancestral state of each site */ tsk_size_t *num_alleles; allele_t *ancestral_state; - allele_t *derived_state; + /* The mutations at site j, in tskit order, are entries + * site_mutations_offset[j] to site_mutations_offset[j + 1] of the + * mutation_node and mutation_derived_state arrays */ + tsk_size_t *site_mutations_offset; + tsk_id_t *mutation_node; + allele_t *mutation_derived_state; + /* Entry j is edge edge_insertion_order[j] (edge_removal_order[j]) of + * the edge table, so that the ranges of a tsk_tree_position_t index + * these arrays directly */ + tsls_edge_t *insertion_edges; + tsls_edge_t *removal_edges; +} tsls_index_t; + +/* The Li and Stephens HMM on a tree sequence. The index is borrowed, and + * must not be changed or freed while the HMM is in use. */ +typedef struct { + int flags; + const tsls_index_t *index; + const tsk_treeseq_t *tree_sequence; + size_t num_nodes; + size_t num_sites; + const double *sites_position; + double sequence_length; /* Input LS model rates */ double likelihood_threshold; int pruning_mode; @@ -154,9 +188,13 @@ int compressed_matrix_alloc( int compressed_matrix_free(compressed_matrix_t *self); size_t compressed_matrix_get_total_values(const compressed_matrix_t *self); -int tsls_hmm_alloc(tsls_hmm_t *self, const tsk_treeseq_t *tree_sequence, - const char **alleles, tsk_size_t num_alleles, const double *recombination_rate, - const double *mismatch_rate, double likelihood_threshold, int flags); +int tsls_index_init(tsls_index_t *self, const tsk_treeseq_t *tree_sequence, + const char **alleles, tsk_size_t num_alleles); +int tsls_index_free(tsls_index_t *self); + +int tsls_hmm_alloc(tsls_hmm_t *self, const tsls_index_t *index, + const double *recombination_rate, const double *mismatch_rate, + double likelihood_threshold, int flags); int tsls_hmm_set_pruning(tsls_hmm_t *self, int mode, double parameter); int tsls_hmm_free(tsls_hmm_t *self); int tsls_hmm_find_path(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, diff --git a/tsls/tests/test_matching.py b/tsls/tests/test_matching.py index b951491..14fa9aa 100644 --- a/tsls/tests/test_matching.py +++ b/tsls/tests/test_matching.py @@ -545,11 +545,13 @@ def test_site_without_mutations(self): def test_shared_tree_sequence(self): ts = two_tree_ts() - tts = tsls.TreeSequence(ts) + tts = tsls.TreeSequence(ts, alleles=ALLELES) assert tts.num_nodes == ts.num_nodes assert tts.num_sites == ts.num_sites m = ts.num_sites - am1 = tsls.Hmm(tts, np.full(m, 0.01), np.full(m, 0.01), alleles=ALLELES) + with pytest.raises(ValueError, match="making the tsls.TreeSequence"): + tsls.Hmm(tts, np.full(m, 0.01), np.full(m, 0.01), alleles=ALLELES) + am1 = tsls.Hmm(tts, np.full(m, 0.01), np.full(m, 0.01)) am2 = tsls.Hmm(ts, np.full(m, 0.01), np.full(m, 0.01), alleles=ALLELES) assert am1.tree_sequence is tts h = np.zeros(m, dtype=np.int8) @@ -586,7 +588,7 @@ def test_bad_lengths(self): tsls.Hmm(ts, np.zeros(1), np.zeros(ts.num_sites)) with pytest.raises(ValueError, match="mismatch"): tsls.Hmm(ts, np.zeros(ts.num_sites), np.zeros(1)) - am, _ = make_matcher(ts, 1e-2, 1e-2, True) + am, _ = make_matcher(single_tree_ts(), 1e-2, 1e-2, True) with pytest.raises(ValueError): am.forward_matrix(np.zeros(1, dtype=np.int8)) h = np.zeros(ts.num_sites, dtype=np.int8) diff --git a/tsls/tsls/matching.py b/tsls/tsls/matching.py index 7d2adab..191104b 100644 --- a/tsls/tsls/matching.py +++ b/tsls/tsls/matching.py @@ -9,17 +9,29 @@ class TreeSequence(_tsls.TreeSequence): """ - A tskit tree sequence in the form used by the HMM. It can be shared - between HMMs, including HMMs in different threads, so that it is only - built once. + A tskit tree sequence in the form used by the HMM, with its alleles + encoded as integers. It can be shared between HMMs, including HMMs in + different threads, so that it is only built once. + + By default the alleles are numbered at each site as tskit does: the + ancestral state is 0 and derived states are numbered from 1 in order of + their first appearance at the site. Every site has at least two alleles, + so that 1 can be matched at a site with no mutations (as in tsinfer, for + an ancestor's focal sites). If ``alleles`` is specified, every state is + instead its index in this list. """ - def __init__(self, ts): + def __init__(self, ts, *, alleles=None): tables = ts.dump_tables() ll_tables = _tsls.LightweightTableCollection(tables.sequence_length) ll_tables.fromdict(tables.asdict()) + if alleles is not None: + alleles = list(alleles) + if len(set(alleles)) != len(alleles): + raise ValueError("Alleles must be unique") self.ts = ts - super().__init__(ll_tables) + self.alleles = alleles + super().__init__(ll_tables, alleles=alleles) class CompressedMatrix: @@ -139,18 +151,14 @@ class Hmm(_tsls.Hmm): Runs the Viterbi, forward and backward algorithms for haplotypes against the specified tree sequence (a tskit.TreeSequence, or a tsls.TreeSequence to share one between HMMs), using the specified per-site recombination - and mismatch probabilities. + and mismatch probabilities. ``alleles`` is passed to tsls.TreeSequence, + and can only be given with a tskit.TreeSequence. The states of the HMM at a given position are the nodes in the tree at that position: nodes with a parent or at least one child. Isolated nodes cannot be copied. - Haplotypes are integer encoded. By default the alleles are numbered at - each site as tskit does: the ancestral state is 0 and derived states are - numbered from 1 in order of their first appearance at the site. Every - site has at least two alleles, so that 1 can be matched at a site with - no mutations (as in tsinfer, for an ancestor's focal sites). If - ``alleles`` is specified, every state is instead its index in this list. + Haplotypes are integer encoded, as described in tsls.TreeSequence. After normalisation at each site, values below likelihood_threshold are set to it. The experimental ``pruning`` option also sets the values of @@ -177,13 +185,10 @@ def __init__( count_distinct=False, ): if not isinstance(ts, TreeSequence): - ts = TreeSequence(ts) + ts = TreeSequence(ts, alleles=alleles) + elif alleles is not None: + raise ValueError("Specify alleles when making the tsls.TreeSequence") kwargs = {"extended_checks": extended_checks, "count_distinct": count_distinct} - if alleles is not None: - alleles = list(alleles) - if len(set(alleles)) != len(alleles): - raise ValueError("Alleles must be unique") - kwargs["alleles"] = alleles if pruning is not None: mode, parameter = pruning if mode not in self.PRUNING_MODES: diff --git a/verification/verification.py b/verification/verification.py index 66c907b..2f9092f 100644 --- a/verification/verification.py +++ b/verification/verification.py @@ -150,12 +150,11 @@ def make_matcher( ): m = ts.num_sites if tree_sequence is None: - tree_sequence = tsls.TreeSequence(ts) + tree_sequence = tsls.TreeSequence(ts, alleles=ALLELES) return tsls.Hmm( tree_sequence, np.full(m, rho), np.full(m, mu), - alleles=ALLELES, weight_by_n=weight_by_n, likelihood_threshold=likelihood_threshold, pruning=pruning, From ebf50a2ac70a7a4be6e77df6b13633bce57625b7 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 13:54:06 +0100 Subject: [PATCH 19/25] Copy the edges into the index in tskit's insertion and removal orders The index now holds the edges in edge_insertion_order and edge_removal_order as 16-byte records (parent, child, and one coordinate), so that the ranges of a tsk_tree_position_t index them directly: tree updates, the traceback and the first-tree seek read them sequentially instead of reading four edge table columns through the order arrays. The seek only needs one coordinate: the edges in its range all have left <= x in insertion order (forward seek) and right > x in removal order (backward seek). Memory: 32 bytes per edge, the same as tsinfer-hmm's copied edge indexes. Paired wall time (paired.py, full set), tsls / tsinfer-hmm: | workload | queries | before (db2415e) | after | |---|---|---|---| | ancestors_g84 | 50 | 1.43 | 0.74 | | ancestors_g1000 | 50 | 1.19 | 0.78 | | ancestors_g1483 | 100 | 1.19 | 0.91 | | ancestors_g1484 | 50 | 1.23 | 1.02 | | ancestors_g1485 | 100 | 1.25 | 0.79 | | samples | 6 | 1.14 | 1.16 | ('before' is before the shared index too, which made most of the difference for g84.) Late ancestors gain most because their seek scans most of the edges, now sequentially. Callgrind's cache simulation is misleading here: it shows last-level read misses on the sample rising from 1.31M to 1.83M, because tskit's tree position scans still read the edge table columns, and the copies add a second stream (about 190k misses per loop). But the simulator doesn't model hardware prefetching, which hides most of the cost of a sequential stream, while the old reads through the order arrays were random. Wall time is the judge from here; the simulator is still useful for locating misses. Dropped: asserts in the seek loops checking the one-coordinate test against the edge table. The Python module builds with asserts on, and they added 25% to the instructions of late ancestors (the range covers most edges). The window tests check the seek against the dense model. Instruction counts: | workload | queries | tsinfer | tsls | ratio | |---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.153 G | 0.518 | | ancestors_g1000 | 20 | 0.621 G | 0.319 G | 0.514 | | ancestors_g1483 | 10 | 1.355 G | 1.060 G | 0.783 | | ancestors_g1484 | 3 | 0.646 G | 0.564 G | 0.873 | | ancestors_g1485 | 5 | 1.131 G | 0.910 G | 0.804 | | samples | 1 (cut at 100k sites) | 1.561 G | 1.522 G | 0.975 | Paths unchanged. --- tsls/lib/hmm.c | 105 +++++++++++++++++++++++++++++++++++-------------- 1 file changed, 76 insertions(+), 29 deletions(-) diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index 5f86049..408478e 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -359,6 +359,56 @@ tsls_index_encode_alleles( return ret; } +/* Copy the edges in tskit's insertion and removal orders, so that the + * tree position's ranges can be read sequentially. */ +static void +tsls_index_copy_edges(tsls_index_t *self) +{ + const tsk_table_collection_t *tables = self->tree_sequence->tables; + const tsk_id_t *restrict insertion_order = tables->indexes.edge_insertion_order; + const tsk_id_t *restrict removal_order = tables->indexes.edge_removal_order; + const tsk_id_t *restrict edge_parent = tables->edges.parent; + const tsk_id_t *restrict edge_child = tables->edges.child; + const double *restrict edge_left = tables->edges.left; + const double *restrict edge_right = tables->edges.right; + tsls_edge_t *restrict insertion_edges = self->insertion_edges; + tsls_edge_t *restrict removal_edges = self->removal_edges; + size_t j; + tsk_id_t e; + + for (j = 0; j < self->num_edges; j++) { + e = insertion_order[j]; + insertion_edges[j].parent = edge_parent[e]; + insertion_edges[j].child = edge_child[e]; + insertion_edges[j].coordinate = edge_right[e]; + e = removal_order[j]; + removal_edges[j].parent = edge_parent[e]; + removal_edges[j].child = edge_child[e]; + removal_edges[j].coordinate = edge_left[e]; + } +} + +/* Return the copied edges for the specified order (one of the ranges of a + * tree position). */ +static inline const tsls_edge_t * +tsls_index_get_edges(const tsls_index_t *self, const tsk_id_t *order) +{ + return order == self->tree_sequence->tables->indexes.edge_insertion_order + ? self->insertion_edges + : self->removal_edges; +} + +/* Returns true if the specified edge in the in range of a tree position + * just seeked from null (see seek_tree_position) covers x. Edges in the + * range all have left <= x in insertion order (a forward seek) and right > x + * in removal order (a backward seek), so one coordinate is enough. */ +static inline bool +seek_edge_covers(const tsk_tree_position_t *tp, const tsls_edge_t *edge, double x) +{ + return tp->direction == TSK_DIR_FORWARD ? edge->coordinate > x + : edge->coordinate <= x; +} + int tsls_index_init(tsls_index_t *self, const tsk_treeseq_t *tree_sequence, const char **alleles, tsk_size_t num_alleles) @@ -381,9 +431,12 @@ tsls_index_init(tsls_index_t *self, const tsk_treeseq_t *tree_sequence, self->mutation_node = malloc((num_mutations + 1) * sizeof(*self->mutation_node)); self->mutation_derived_state = malloc((num_mutations + 1) * sizeof(*self->mutation_derived_state)); + self->insertion_edges = malloc((self->num_edges + 1) * sizeof(*self->insertion_edges)); + self->removal_edges = malloc((self->num_edges + 1) * sizeof(*self->removal_edges)); if (self->num_alleles == NULL || self->ancestral_state == NULL || self->site_mutations_offset == NULL || self->mutation_node == NULL - || self->mutation_derived_state == NULL) { + || self->mutation_derived_state == NULL || self->insertion_edges == NULL + || self->removal_edges == NULL) { ret = TSLS_ERR_NO_MEMORY; goto out; } @@ -391,6 +444,7 @@ tsls_index_init(tsls_index_t *self, const tsk_treeseq_t *tree_sequence, if (ret != 0) { goto out; } + tsls_index_copy_edges(self); out: return ret; } @@ -851,25 +905,23 @@ tsls_hmm_insert_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, bool update_likel static void tsls_hmm_update_tree(tsls_hmm_t *self, const tsk_tree_position_t *tp) { - const tsk_id_t *restrict edge_parent = self->tree_sequence->tables->edges.parent; - const tsk_id_t *restrict edge_child = self->tree_sequence->tables->edges.child; + const tsls_edge_t *restrict out = tsls_index_get_edges(self->index, tp->out.order); + const tsls_edge_t *restrict in = tsls_index_get_edges(self->index, tp->in.order); const int direction = tp->direction; tsk_id_t *restrict isolated = self->likelihood_nodes_tmp; double *restrict L = self->likelihood; - tsk_id_t j, e, u; + tsk_id_t j, u; int k, num_left, num_isolated = 0; for (j = tp->out.start; j != tp->out.stop; j += direction) { - e = tp->out.order[j]; tsls_hmm_remove_edge( - self, edge_parent[e], edge_child[e], true, isolated, &num_isolated); + self, out[j].parent, out[j].child, true, isolated, &num_isolated); } for (j = tp->out.start; j != tp->out.stop; j += direction) { - tsls_hmm_reset_likelihood_cache(self, edge_parent[tp->out.order[j]]); + tsls_hmm_reset_likelihood_cache(self, out[j].parent); } for (j = tp->in.start; j != tp->in.stop; j += direction) { - e = tp->in.order[j]; - tsls_hmm_insert_edge(self, edge_parent[e], edge_child[e], true); + tsls_hmm_insert_edge(self, in[j].parent, in[j].child, true); } /* Nodes that are still isolated have left the tree. Nodes that were * isolated and then inserted again keep their likelihoods. We mark the @@ -897,11 +949,8 @@ tsls_hmm_init_sweep( { int ret = 0; const size_t N = self->num_nodes; - const tsk_id_t *restrict edge_parent = self->tree_sequence->tables->edges.parent; - const tsk_id_t *restrict edge_child = self->tree_sequence->tables->edges.child; - const double *restrict edge_left = self->tree_sequence->tables->edges.left; - const double *restrict edge_right = self->tree_sequence->tables->edges.right; - tsk_id_t u, j, e; + const tsls_edge_t *restrict in; + tsk_id_t u, j; size_t k; memset(self->parent, 0xff, N * sizeof(*self->parent)); @@ -920,10 +969,10 @@ tsls_hmm_init_sweep( if (ret != 0) { goto out; } + in = tsls_index_get_edges(self->index, tp->in.order); for (j = tp->in.start; j != tp->in.stop; j += tp->direction) { - e = tp->in.order[j]; - if (edge_left[e] <= position && position < edge_right[e]) { - tsls_hmm_insert_edge(self, edge_parent[e], edge_child[e], false); + if (seek_edge_covers(tp, &in[j], position)) { + tsls_hmm_insert_edge(self, in[j].parent, in[j].child, false); } } self->num_likelihood_nodes = 0; @@ -1544,7 +1593,7 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, { int ret = 0; tsk_tree_position_t tp; - tsk_id_t site, j, e, u, max_likelihood_node; + tsk_id_t site, j, u, max_likelihood_node; double x; bool valid; tsk_id_t *restrict parent = self->parent; @@ -1555,11 +1604,8 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, * whose allelic state is the site's ancestral state, so the walk to find * a node's allele doesn't need to check for the root. */ const tsk_id_t virtual_root = (tsk_id_t) self->num_nodes; - const tsk_edge_table_t *edges = &self->tree_sequence->tables->edges; - const tsk_id_t *restrict edge_parent = edges->parent; - const tsk_id_t *restrict edge_child = edges->child; - const double *restrict edge_left = edges->left; - const double *restrict edge_right = edges->right; + const tsls_edge_t *restrict in; + const tsls_edge_t *restrict out; const double start_pos = start == 0 ? 0 : sites_position[start]; const double end_pos = end == (tsk_id_t) self->num_sites ? self->sequence_length : sites_position[end]; @@ -1588,10 +1634,10 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, if (ret != 0) { goto out; } + in = tsls_index_get_edges(self->index, tp.in.order); for (j = tp.in.start; j != tp.in.stop; j += tp.direction) { - e = tp.in.order[j]; - if (edge_left[e] <= x && x < edge_right[e]) { - parent[edge_child[e]] = edge_parent[e]; + if (seek_edge_covers(&tp, &in[j], x)) { + parent[in[j].child] = in[j].parent; } } for (site = end - 1; site >= start; site--) { @@ -1599,12 +1645,13 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, valid = tsk_tree_position_prev(&tp); assert(valid); (void) valid; + out = tsls_index_get_edges(self->index, tp.out.order); + in = tsls_index_get_edges(self->index, tp.in.order); for (j = tp.out.start; j != tp.out.stop; j--) { - parent[edge_child[tp.out.order[j]]] = virtual_root; + parent[out[j].child] = virtual_root; } for (j = tp.in.start; j != tp.in.stop; j--) { - e = tp.in.order[j]; - parent[edge_child[e]] = edge_parent[e]; + parent[in[j].child] = in[j].parent; } } allelic_state[virtual_root] = self->index->ancestral_state[site]; From 9c7a978f2f9843d9f4ace56ed16154f00f5c1096 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 14:38:57 +0100 Subject: [PATCH 20/25] Record wall times after the shared index, and what didn't help Documents the paired wall times after the shared index (ancestors now 0.74-1.02 of tsinfer-hmm, samples 1.07-1.16), two reverted experiments (removing the per-edge cache clear, and marking roots at each site again: no measurable change on samples), and a perf profile locating the samples' remaining gap in the likelihood cache walks, which execute about the same instructions as tsinfer-hmm's but miss the cache more. --- benchmarking/README.md | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/benchmarking/README.md b/benchmarking/README.md index b7f49f9..4d560a3 100644 --- a/benchmarking/README.md +++ b/benchmarking/README.md @@ -518,3 +518,32 @@ Two things account for this: mutations from `tsk_mutation_t` structs. The extra L1 write misses come from two of the optimisations, which write per-edge marks into `allelic_state` and `likelihood_cache`. + +After sharing a `tsls_index_t` between HMMs (allele encoding, compact +mutations, and the edges copied in tskit's insertion and removal orders), +paired wall times, tsls / tsinfer-hmm: + +| workload | queries | ratio | +|---|---|---| +| ancestors, group 84 | 50 | 0.74 | +| ancestors, group 1000 | 50 | 0.78 | +| ancestors, group 1483 | 100 | 0.91 | +| ancestors, group 1484 | 50 | 1.02 | +| ancestors, group 1485 | 100 | 0.79 | +| samples | 6 | 1.07-1.16 | + +Paired ratios for the samples vary by a few percent between runs of the +same build. Neither per-edge write added by the instruction count work +explains the remaining gap: without the cache clear in `remove_edge`, the +samples gave 1.07 and 1.10 (against 1.07 and 1.10 with it), and with the +roots marked at each site instead of throughout the sweep, 1.11 and 1.11. + +A sampling profile (perf, 2 whole samples) puts tsls's matching at 1.21 +times tsinfer-hmm's samples. The allele walk costs the same in both. The +difference is in the likelihood cache walks (finding inherited values, +filling and resetting the cache, inserting mutation likelihoods), which +execute about the same instructions as tsinfer-hmm's but take 30-80% +longer, and in checking for nodes leaving the tree. Each step of these +walks reads several per-node arrays (`L`, `L_cache`, the flag arrays, +`parent`); with 433k nodes they total more than the 16 MB L3 cache on the +test machine, and tsls uses more of them than tsinfer-hmm. From 9806f5596df380fc43e0f1956e602ebef03521e9 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 16:05:54 +0100 Subject: [PATCH 21/25] Back the large per-node and edge arrays with transparent huge pages The likelihood walks jump around per-node arrays of several MB each (433k nodes for the chr20p ancestors), so with 4 KB pages most steps miss the TLB. Large arrays (>= 1 MiB: the per-node arrays and the index's edge copies) are now mapped directly, aligned to 2 MiB, with madvise(MADV_HUGEPAGE) before first use (Linux only; elsewhere and for small arrays, malloc). They are mapped directly because glibc reuses memory already faulted in with small pages for large requests in a long-running process (its mmap threshold grows after large frees), and madvise doesn't change pages that are already there: with posix_memalign + madvise, no huge pages were used at all. Paired wall time, tsls / tsinfer-hmm (paired.py), two rounds run alternately with the previous build (U) to share machine conditions: | workload | U round 1 | U round 2 | this, round 1 | this, round 2 | |---|---|---|---|---| | ancestors_g84 | 0.75 | 0.75 | 0.72 | 0.75 | | ancestors_g1000 | 0.76 | 0.75 | 0.76 | 0.74 | | ancestors_g1483 | 0.92 | 0.91 | 0.91 | 0.80 | | ancestors_g1484 | 0.99 | 1.00 | 0.96 | 0.79 | | ancestors_g1485 | 1.06 | 1.04 | 1.00 | 1.00 | | samples | 0.99 | 1.11 | 0.87 | 0.81 | tsls's own times also vary less between runs (samples 16.99 and 19.24 s before; 14.67 and 14.22 s now), while tsinfer-hmm's were steady throughout. Run-to-run variation of tsls of up to 20% had made earlier single comparisons unreliable; it fits page placement varying between processes. Dead ends (not committed), each run twice on the full paired set, alternating with the then-current build: - One 24-byte record per node {likelihood, likelihood_cache, parent, allelic_state, recombination_required, its cache}, so that each step of a walk reads one record: samples 1.52 and 1.52, ancestors 0.80-1.64, all worse. Indexing 24-byte records also costs an address computation per step: instructions on the profiling sample rose from 1.52G to 1.82G (the allele walk alone 538M -> 804M). - Split records, an 8-byte {parent, allelic_state, flags} and a 16-byte {likelihood, likelihood_cache}: no consistent difference from separate arrays (samples 1.16 and 0.94 against 1.17 and 1.12; ancestors within noise). Paths unchanged. Instruction counts are unaffected. --- tsls/lib/hmm.c | 153 ++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 119 insertions(+), 34 deletions(-) diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index 408478e..4424347 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -25,6 +25,11 @@ * is in the tree over the whole interval between them. */ +#if defined(__linux__) +/* For posix_memalign and madvise */ +#define _DEFAULT_SOURCE +#endif + #include "tsls.h" #include @@ -32,6 +37,9 @@ #include #include #include +#if defined(__linux__) +#include +#endif /* For the innermost helpers, which GCC doesn't always inline. Inlining * tsls_hmm_coalesce_likelihoods also specialises it for its constant @@ -42,6 +50,81 @@ #define ALWAYS_INLINE inline #endif +/* Allocate an array that is accessed at random, such as the per-node arrays + * walked up the tree. On Linux, large arrays are mapped directly, aligned + * to 2 MiB, and backed by transparent huge pages where available + * (MADV_HUGEPAGE), which greatly reduces TLB misses in the walks. They are + * mapped directly because memory from malloc may already have been touched + * with small pages, which madvise doesn't change. Free with tsls_large_free, + * with the same size. */ +#define HUGE_PAGE_SIZE ((size_t) 2 * 1024 * 1024) + +static inline bool +tsls_large_alloc_uses_mmap(size_t size) +{ +#if defined(__linux__) && defined(MADV_HUGEPAGE) + return size >= HUGE_PAGE_SIZE / 2; +#else + (void) size; + return false; +#endif +} + +static void * +tsls_large_alloc(size_t size) +{ +#if defined(__linux__) && defined(MADV_HUGEPAGE) + size_t rounded, head, tail; + char *p, *aligned; + + if (tsls_large_alloc_uses_mmap(size)) { + rounded = (size + HUGE_PAGE_SIZE - 1) & ~(HUGE_PAGE_SIZE - 1); + p = mmap(NULL, rounded + HUGE_PAGE_SIZE, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (p == MAP_FAILED) { + return NULL; + } + /* Trim the mapping to a range aligned to the huge page size */ + aligned = (char *) (((uintptr_t) p + HUGE_PAGE_SIZE - 1) + & ~((uintptr_t) HUGE_PAGE_SIZE - 1)); + head = (size_t) (aligned - p); + tail = HUGE_PAGE_SIZE - head; + if (head > 0) { + (void) munmap(p, head); + } + if (tail > 0) { + (void) munmap(aligned + rounded, tail); + } + /* Huge pages are only advice, so failure doesn't matter */ + (void) madvise(aligned, rounded, MADV_HUGEPAGE); + return aligned; + } +#endif + return malloc(size); +} + +static void +tsls_large_free(void *p, size_t size) +{ + if (p == NULL) { + return; + } +#if defined(__linux__) && defined(MADV_HUGEPAGE) + if (tsls_large_alloc_uses_mmap(size)) { + (void) munmap(p, (size + HUGE_PAGE_SIZE - 1) & ~(HUGE_PAGE_SIZE - 1)); + return; + } +#endif + free(p); +} + +/* Free an array of N + 1 elements allocated with tsls_large_alloc */ +#define tsls_large_free_array(array, n) \ + do { \ + tsls_large_free(array, ((n) + 1) * sizeof(*(array))); \ + array = NULL; \ + } while (0) + #define MODE_VITERBI 0 #define MODE_FORWARD 1 #define MODE_BACKWARD 2 @@ -431,8 +514,10 @@ tsls_index_init(tsls_index_t *self, const tsk_treeseq_t *tree_sequence, self->mutation_node = malloc((num_mutations + 1) * sizeof(*self->mutation_node)); self->mutation_derived_state = malloc((num_mutations + 1) * sizeof(*self->mutation_derived_state)); - self->insertion_edges = malloc((self->num_edges + 1) * sizeof(*self->insertion_edges)); - self->removal_edges = malloc((self->num_edges + 1) * sizeof(*self->removal_edges)); + self->insertion_edges + = tsls_large_alloc((self->num_edges + 1) * sizeof(*self->insertion_edges)); + self->removal_edges + = tsls_large_alloc((self->num_edges + 1) * sizeof(*self->removal_edges)); if (self->num_alleles == NULL || self->ancestral_state == NULL || self->site_mutations_offset == NULL || self->mutation_node == NULL || self->mutation_derived_state == NULL || self->insertion_edges == NULL @@ -457,8 +542,8 @@ tsls_index_free(tsls_index_t *self) tsi_safe_free(self->site_mutations_offset); tsi_safe_free(self->mutation_node); tsi_safe_free(self->mutation_derived_state); - tsi_safe_free(self->insertion_edges); - tsi_safe_free(self->removal_edges); + tsls_large_free_array(self->insertion_edges, self->num_edges); + tsls_large_free_array(self->removal_edges, self->num_edges); return 0; } @@ -567,23 +652,23 @@ tsls_hmm_alloc(tsls_hmm_t *self, const tsls_index_t *index, self->traceback = calloc(self->num_sites + 1, sizeof(node_state_list_t)); self->max_likelihood_node = malloc((self->num_sites + 1) * sizeof(tsk_id_t)); - self->parent = malloc((N + 1) * sizeof(*self->parent)); - self->left_child = malloc((N + 1) * sizeof(*self->left_child)); - self->right_child = malloc((N + 1) * sizeof(*self->right_child)); - self->left_sib = malloc((N + 1) * sizeof(*self->left_sib)); - self->right_sib = malloc((N + 1) * sizeof(*self->right_sib)); - self->num_below = malloc((N + 1) * sizeof(*self->num_below)); - self->num_inheriting = malloc((N + 1) * sizeof(*self->num_inheriting)); - self->likelihood = malloc((N + 1) * sizeof(*self->likelihood)); - self->allelic_state = malloc((N + 1) * sizeof(*self->allelic_state)); + self->parent = tsls_large_alloc((N + 1) * sizeof(*self->parent)); + self->left_child = tsls_large_alloc((N + 1) * sizeof(*self->left_child)); + self->right_child = tsls_large_alloc((N + 1) * sizeof(*self->right_child)); + self->left_sib = tsls_large_alloc((N + 1) * sizeof(*self->left_sib)); + self->right_sib = tsls_large_alloc((N + 1) * sizeof(*self->right_sib)); + self->num_below = tsls_large_alloc((N + 1) * sizeof(*self->num_below)); + self->num_inheriting = tsls_large_alloc((N + 1) * sizeof(*self->num_inheriting)); + self->likelihood = tsls_large_alloc((N + 1) * sizeof(*self->likelihood)); + self->allelic_state = tsls_large_alloc((N + 1) * sizeof(*self->allelic_state)); self->recombination_required - = malloc((N + 1) * sizeof(*self->recombination_required)); - self->likelihood_cache = malloc((N + 1) * sizeof(*self->likelihood_cache)); + = tsls_large_alloc((N + 1) * sizeof(*self->recombination_required)); + self->likelihood_cache = tsls_large_alloc((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)); + = tsls_large_alloc((N + 1) * sizeof(*self->recombination_required_cache)); + self->likelihood_nodes = tsls_large_alloc((N + 1) * sizeof(*self->likelihood_nodes)); + self->likelihood_nodes_tmp = tsls_large_alloc((N + 1) * sizeof(*self->likelihood_nodes_tmp)); + self->values_tmp = tsls_large_alloc((N + 1) * sizeof(*self->values_tmp)); if (self->recombination_rate == NULL || self->mismatch_rate == NULL || self->traceback == NULL @@ -625,21 +710,21 @@ tsls_hmm_free(tsls_hmm_t *self) { tsi_safe_free(self->recombination_rate); tsi_safe_free(self->mismatch_rate); - tsi_safe_free(self->parent); - tsi_safe_free(self->left_child); - tsi_safe_free(self->right_child); - tsi_safe_free(self->left_sib); - tsi_safe_free(self->right_sib); - tsi_safe_free(self->num_below); - tsi_safe_free(self->num_inheriting); - 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); - tsi_safe_free(self->allelic_state); + tsls_large_free_array(self->parent, self->num_nodes); + tsls_large_free_array(self->left_child, self->num_nodes); + tsls_large_free_array(self->right_child, self->num_nodes); + tsls_large_free_array(self->left_sib, self->num_nodes); + tsls_large_free_array(self->right_sib, self->num_nodes); + tsls_large_free_array(self->num_below, self->num_nodes); + tsls_large_free_array(self->num_inheriting, self->num_nodes); + tsls_large_free_array(self->recombination_required, self->num_nodes); + tsls_large_free_array(self->likelihood, self->num_nodes); + tsls_large_free_array(self->likelihood_cache, self->num_nodes); + tsls_large_free_array(self->recombination_required_cache, self->num_nodes); + tsls_large_free_array(self->likelihood_nodes, self->num_nodes); + tsls_large_free_array(self->likelihood_nodes_tmp, self->num_nodes); + tsls_large_free_array(self->values_tmp, self->num_nodes); + tsls_large_free_array(self->allelic_state, self->num_nodes); tsi_safe_free(self->max_likelihood_node); tsi_safe_free(self->traceback); tsi_blkalloc_free(&self->traceback_allocator); From 5b5865138d0b2fad57c9710b09cbebc8e226bc70 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 16:09:26 +0100 Subject: [PATCH 22/25] Update the benchmarking README with the wall-time results Replaces the interim wall-time notes with: how paired.py measures, the current result (tsls faster than tsinfer-hmm on samples by 13-19% and on most ancestor groups, level on group 1485), how each change moved it (after the refactor, + shared index, + huge pages), why single runs and callgrind's cache simulation mislead on this machine, and the node layouts that didn't help. --- benchmarking/README.md | 160 ++++++++++++++++++++++++++--------------- 1 file changed, 101 insertions(+), 59 deletions(-) diff --git a/benchmarking/README.md b/benchmarking/README.md index 4d560a3..d60c226 100644 --- a/benchmarking/README.md +++ b/benchmarking/README.md @@ -487,63 +487,105 @@ its traceback sizes are unchanged. ### Wall time after the refactor -Fewer instructions did not mean less time. `paired.py` times each query -three times with each engine, alternating, and keeps the best time for -each (on an idle machine; the ratio is total tsls time over total -tsinfer-hmm time, including building the matcher for each query, as -`engines.py` does): +Fewer instructions did not at first mean less time, so wall time was then +optimised directly. `paired.py` matches each query three times with each +engine, alternating between them, keeps the best time for each, and +reports total tsls time over total tsinfer-hmm time for each workload +(the per-query median ratio is also printed). The times include building +the matcher for each query, as `engines.py` does. `paired.py --quick` +runs a smaller set, but it is too noisy on this VM to compare builds (see +below). -| workload | queries | tsinfer-hmm (s) | tsls (s) | ratio | -|---|---|---|---|---| -| ancestors, group 84 | 50 | 0.21 | 0.30 | 1.43 | -| ancestors, group 1000 | 50 | 0.46 | 0.55 | 1.19 | -| ancestors, group 1483 | 100 | 4.76 | 5.67 | 1.19 | -| ancestors, group 1484 | 50 | 5.23 | 6.42 | 1.23 | -| ancestors, group 1485 | 100 | 12.98 | 16.20 | 1.25 | -| samples | 6 | 17.69 | 20.22 | 1.14 | - -Two things account for this: - -- Building a `tsls.Hmm` takes 4-8 ms against 1.5 ms for tsinfer-hmm's - `AncestorMatcher`, because it encodes the alleles of every site from the - tables each time. This dominates for early ancestors. -- `find_path` itself is slower (1.10 on group 1485, 1.18 on samples, timed - separately) despite executing fewer instructions, because it misses the - cache more. Callgrind's cache simulation on the profiling sample: tsls - has 2% fewer instructions but 25% more L1 read misses, 2.2 times the L1 - write misses and 35% more last-level read misses than tsinfer-hmm. The - last-level misses are mostly in tskit's tree position scans, which read - edge coordinates through the index arrays (`left[order[j]]`) where - tsinfer-hmm scans a contiguous array of sorted edges, and in reading - mutations from `tsk_mutation_t` structs. The extra L1 write misses come - from two of the optimisations, which write per-edge marks into - `allelic_state` and `likelihood_cache`. - -After sharing a `tsls_index_t` between HMMs (allele encoding, compact -mutations, and the edges copied in tskit's insertion and removal orders), -paired wall times, tsls / tsinfer-hmm: - -| workload | queries | ratio | -|---|---|---| -| ancestors, group 84 | 50 | 0.74 | -| ancestors, group 1000 | 50 | 0.78 | -| ancestors, group 1483 | 100 | 0.91 | -| ancestors, group 1484 | 50 | 1.02 | -| ancestors, group 1485 | 100 | 0.79 | -| samples | 6 | 1.07-1.16 | - -Paired ratios for the samples vary by a few percent between runs of the -same build. Neither per-edge write added by the instruction count work -explains the remaining gap: without the cache clear in `remove_edge`, the -samples gave 1.07 and 1.10 (against 1.07 and 1.10 with it), and with the -roots marked at each site instead of throughout the sweep, 1.11 and 1.11. - -A sampling profile (perf, 2 whole samples) puts tsls's matching at 1.21 -times tsinfer-hmm's samples. The allele walk costs the same in both. The -difference is in the likelihood cache walks (finding inherited values, -filling and resetting the cache, inserting mutation likelihoods), which -execute about the same instructions as tsinfer-hmm's but take 30-80% -longer, and in checking for nodes leaving the tree. Each step of these -walks reads several per-node arrays (`L`, `L_cache`, the flag arrays, -`parent`); with 433k nodes they total more than the 16 MB L3 cache on the -test machine, and tsls uses more of them than tsinfer-hmm. +``` +uv run python paired.py +``` + +#### Current result + +Two full runs of the current build (`9806f55`), tsls / tsinfer-hmm: + +| workload | queries | round 1 | round 2 | +|---|---|---|---| +| ancestors, group 84 | 50 | 0.72 | 0.75 | +| ancestors, group 1000 | 50 | 0.76 | 0.74 | +| ancestors, group 1483 | 100 | 0.91 | 0.80 | +| ancestors, group 1484 | 50 | 0.96 | 0.79 | +| ancestors, group 1485 | 100 | 1.00 | 1.00 | +| samples | 6 | 0.87 | 0.81 | + +tsls is now faster than tsinfer-hmm on samples (by 13-19%) and on most +ancestor groups, and level on group 1485, with the same paths and +traceback sizes. + +#### How it got there + +| workload | after the refactor | + shared index | + huge pages | +|---|---|---|---| +| ancestors, group 84 | 1.43 | 0.73-0.77 | 0.72-0.75 | +| ancestors, group 1000 | 1.19 | 0.75-0.76 | 0.74-0.76 | +| ancestors, group 1483 | 1.19 | 0.88-0.92 | 0.80-0.91 | +| ancestors, group 1484 | 1.23 | 0.99-1.00 | 0.79-0.96 | +| ancestors, group 1485 | 1.25 | 0.97-1.06 | 1.00 | +| samples | 1.14 | 0.99-1.17 | 0.81-0.87 | + +(Ranges are over four runs of the shared index build and two of the huge +pages build, run alternately with other builds.) + +- **After the refactor**, tsls was slower despite executing fewer + instructions. Building a `tsls.Hmm` took 4-8 ms against 1.5 ms for + tsinfer-hmm's `AncestorMatcher`, because each one encoded the alleles + of every site from the tables, which dominated for early ancestors. + And `find_path` missed the cache more: callgrind's cache simulation on + the profiling sample showed 2% fewer instructions but 25% more L1 read + misses and 35% more last-level read misses than tsinfer-hmm, mostly from + reading edges through tskit's order arrays and mutations from + `tsk_mutation_t` structs. +- **Shared index** (`92d471b`, `ebf50a2`): a `tsls_index_t`, built once + per tree sequence and shared by HMMs, holds the allele encoding, compact + mutation arrays, and the edges copied in tskit's insertion and removal + orders, which the ranges of a `tsk_tree_position_t` index directly. + Late ancestors gain most, because finding their first tree scans most of + the edges, now sequentially. +- **Huge pages** (`9806f55`): the likelihood walks jump around per-node + arrays of several MB each (433k nodes), so with 4 KB pages most steps + miss the TLB. The large arrays are now mapped directly, aligned to + 2 MiB, with `madvise(MADV_HUGEPAGE)` (Linux). They have to be mapped + directly: glibc reuses memory already faulted in with small pages, which + `madvise` doesn't change, so `posix_memalign` + `madvise` gave no huge + pages at all. + +#### Measuring on this machine + +- Single runs mislead. tsls's time for the same build varied by up to 20% + between processes (6 samples: 16.4 to 19.9 s) while tsinfer-hmm's was + steady, which fits physical page placement varying between runs; with + huge pages tsls's times are steadier (14.2 and 14.7 s). Another session + on the machine may also compete for memory bandwidth at times. Builds + were therefore compared with full runs repeated in alternation (build A, + build B, A, B), and a single `--quick` run once showed 1.36 for group + 1485 where full runs of the same build gave 0.85-0.92. +- Callgrind's cache simulation is useful for finding misses but not for + judging changes: it showed the edge copies increasing last-level misses + (1.31M to 1.83M on the profiling sample, since tskit's scans still read + the edge table and the copies add a second stream), yet they were much + faster, because the simulator doesn't model the hardware prefetching + that hides the cost of a sequential stream. +- Instruction counts and wall time can move in opposite directions: + reverting two changes that reduced instructions by adding per-edge + writes (the cache clear in `remove_edge`, and keeping roots marked + through the sweep) made no measurable difference to the samples. + +#### Node layouts that didn't help + +A `perf` profile of the shared index build put the samples' remaining +time in the likelihood cache walks, which read several per-node arrays at +each step. Two layouts that put these fields together were tried, each in +two full runs alternating with the then-current build: + +- One 24-byte record per node (likelihood, its cache, parent, allelic + state, recombination flag and its cache): worse everywhere (samples + 1.52 and 1.52; ancestors up to 1.64). Indexing 24-byte records also + costs an address computation at every step of the walks: instructions + on the profiling sample rose from 1.52G to 1.82G. +- An 8-byte record {parent, allelic state, flags} and a 16-byte record + {likelihood, its cache}: no consistent difference from separate arrays. From 832fc5f00f91f5676325173351c41c1e8d71d209 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 16:29:55 +0100 Subject: [PATCH 23/25] Keep each node's walk state and likelihood in two records The per-node arrays the walks read are now two record arrays: tsls_node_t, 8 bytes {parent, allelic_state, recombination_required, its cache}, and tsls_likelihood_t, 16 bytes {value, cache}. The allele walk reads one record per step instead of two arrays, and the likelihood cache walks read two instead of four. Both arrays are backed by huge pages as before; the tree's child and sibling links and the forward/backward counts stay as separate arrays. Without huge pages this layout made no consistent difference (see 9806f55). With them it is faster where most of the matching time goes. Paired wall time, tsls / tsinfer-hmm, two full runs of each build in alternation (this run's H rounds varied more than in 9806f55): | workload | separate arrays (9806f55) | two records | |---|---|---| | ancestors_g84 | 0.76, 0.76 | 0.80, 0.81 | | ancestors_g1000 | 0.75, 0.76 | 0.82, 0.82 | | ancestors_g1483 | 0.80, 0.96 | 0.78, 0.78 | | ancestors_g1484 | 0.81, 0.92 | 0.84, 0.82 | | ancestors_g1485 | 0.99, 0.95 | 0.90, 0.87 | | samples | 1.11, 0.95 | 0.79, 0.82 | The early groups (4-8 ms per query) are slightly slower; the late groups and the samples (0.1-3 s per query), which dominate tsinfer's matching time, are faster and steadier. The cost is instructions: the profiling sample goes from 1.52G to 1.76G (tsinfer-hmm 1.56G), so tsls is no longer at instruction-count parity on samples. GCC forms a pointer to the record at each step of the allele walk (an extra lea), which is 538M -> 669M instructions there, but the step still costs one memory access. | workload | queries | tsinfer | tsls | ratio | |---|---|---|---|---| | ancestors_g84 | 20 | 0.296 G | 0.157 G | 0.532 | | ancestors_g1000 | 20 | 0.621 G | 0.346 G | 0.557 | | ancestors_g1483 | 10 | 1.355 G | 1.155 G | 0.852 | | ancestors_g1484 | 3 | 0.646 G | 0.629 G | 0.974 | | ancestors_g1485 | 5 | 1.131 G | 0.968 G | 0.856 | | samples | 1 (cut at 100k sites) | 1.561 G | 1.764 G | 1.130 | Paths unchanged. --- tsls/lib/hmm.c | 376 ++++++++++++++++++++++-------------------------- tsls/lib/tsls.h | 35 +++-- 2 files changed, 200 insertions(+), 211 deletions(-) diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c index 4424347..c130d7e 100644 --- a/tsls/lib/hmm.c +++ b/tsls/lib/hmm.c @@ -130,7 +130,7 @@ tsls_large_free(void *p, size_t size) #define MODE_BACKWARD 2 static inline void -remove_edge(tsk_id_t p, tsk_id_t c, tsk_id_t *restrict parent, +remove_edge(tsk_id_t p, tsk_id_t c, tsls_node_t *restrict tree, tsk_id_t *restrict left_child, tsk_id_t *restrict right_child, tsk_id_t *restrict left_sib, tsk_id_t *restrict right_sib) { @@ -147,19 +147,19 @@ remove_edge(tsk_id_t p, tsk_id_t c, tsk_id_t *restrict parent, } else { left_sib[rsib] = lsib; } - parent[c] = NULL_NODE; + tree[c].parent = NULL_NODE; left_sib[c] = NULL_NODE; right_sib[c] = NULL_NODE; } static inline void -insert_edge(tsk_id_t p, tsk_id_t c, tsk_id_t *restrict parent, +insert_edge(tsk_id_t p, tsk_id_t c, tsls_node_t *restrict tree, tsk_id_t *restrict left_child, tsk_id_t *restrict right_child, tsk_id_t *restrict left_sib, tsk_id_t *restrict right_sib) { const tsk_id_t u = right_child[p]; - parent[c] = p; + tree[c].parent = p; if (u == NULL_NODE) { left_child[p] = c; left_sib[c] = NULL_NODE; @@ -177,22 +177,22 @@ insert_edge(tsk_id_t p, tsk_id_t c, tsk_id_t *restrict parent, * site. insert_edge and remove_edge keep this up to date. */ #define ROOT_ALLELE (-2) -/* Returns the allelic state of node u, given the allelic_state array in - * which the nodes with mutations at the current site are set and the roots +/* Returns the allelic state of node u, given a tree in which the nodes with + * mutations at the current site have their allelic states set and the roots * are marked with ROOT_ALLELE, so the walk up the tree doesn't need to * check for the root. This is the innermost loop of the HMM, so there is * deliberately no assert here (the Python module is compiled with asserts * on); the marks are checked in tsls_hmm_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, allele_t ancestral_state) +get_marked_allelic_state( + const tsls_node_t *restrict tree, tsk_id_t u, allele_t ancestral_state) { allele_t state; - while (allelic_state[u] == TSK_NULL) { - u = parent[u]; + while (tree[u].allelic_state == TSK_NULL) { + u = tree[u].parent; } - state = allelic_state[u]; + state = tree[u].allelic_state; return likely(state != ROOT_ALLELE) ? state : ancestral_state; } @@ -284,12 +284,13 @@ aligned_size(size_t size) return (size + 7) & ~((size_t) 7); } -/* Store the values of the specified nodes at this site, and if R is not NULL - * their recombination_required flags. */ +/* Store the likelihoods of the specified nodes at this site, and if + * with_flags is true their recombination_required flags. */ static int WARN_UNUSED compressed_matrix_store_site(compressed_matrix_t *self, tsk_id_t site, double normalisation_factor, int num_values, const tsk_id_t *restrict nodes, - const double *restrict L, const int8_t *restrict R) + const tsls_likelihood_t *restrict L, const tsls_node_t *restrict tree, + bool with_flags) { int ret = 0; int j; @@ -297,7 +298,7 @@ compressed_matrix_store_site(compressed_matrix_t *self, tsk_id_t site, double *restrict values_out; int8_t *restrict flags_out = NULL; - if (R != NULL) { + if (with_flags) { flags_out = tsi_blkalloc_get( &self->allocator, aligned_size((size_t) num_values * sizeof(int8_t))); if (flags_out == NULL) { @@ -305,7 +306,7 @@ compressed_matrix_store_site(compressed_matrix_t *self, tsk_id_t site, goto out; } for (j = 0; j < num_values; j++) { - flags_out[j] = R[nodes[j]]; + flags_out[j] = tree[nodes[j]].recombination_required; } } @@ -319,7 +320,7 @@ compressed_matrix_store_site(compressed_matrix_t *self, tsk_id_t site, } for (j = 0; j < num_values; j++) { nodes_out[j] = nodes[j]; - values_out[j] = L[nodes[j]]; + values_out[j] = L[nodes[j]].value; } self->normalisation_factor[site] = normalisation_factor; self->num_values[site] = (tsk_size_t) num_values; @@ -555,7 +556,7 @@ tsls_index_free(tsls_index_t *self) static inline bool tsls_hmm_in_tree(const tsls_hmm_t *self, tsk_id_t u) { - return self->parent[u] != NULL_NODE || self->left_child[u] != NULL_NODE; + return self->tree[u].parent != NULL_NODE || self->left_child[u] != NULL_NODE; } static void @@ -567,24 +568,24 @@ tsls_hmm_check_state(const tsls_hmm_t *self) for (j = 0; j < self->num_likelihood_nodes; j++) { u = self->likelihood_nodes[j]; - assert(self->likelihood[u] >= 0); + assert(self->likelihood[u].value >= 0); } for (u = 0; u < (tsk_id_t) self->num_nodes; u++) { - if (self->likelihood[u] >= 0) { + if (self->likelihood[u].value >= 0) { num_likelihoods++; } if (tsls_hmm_in_tree(self, u)) { - assert(self->likelihood[u] != NOT_IN_TREE); - if (self->parent[u] == NULL_NODE) { + assert(self->likelihood[u].value != NOT_IN_TREE); + if (self->tree[u].parent == NULL_NODE) { /* Roots always have a likelihood */ - assert(self->likelihood[u] >= 0); + assert(self->likelihood[u].value >= 0); } } else { - assert(self->likelihood[u] == NOT_IN_TREE); + assert(self->likelihood[u].value == NOT_IN_TREE); } - assert(self->allelic_state[u] - == (self->parent[u] == NULL_NODE ? ROOT_ALLELE : TSK_NULL)); - assert(self->likelihood_cache[u] == CACHE_UNSET); + assert(self->tree[u].allelic_state + == (self->tree[u].parent == NULL_NODE ? ROOT_ALLELE : TSK_NULL)); + assert(self->likelihood[u].cache == CACHE_UNSET); } assert(num_likelihoods == self->num_likelihood_nodes); (void) num_likelihoods; @@ -606,14 +607,14 @@ tsls_hmm_print_state(tsls_hmm_t *self, FILE *out) fprintf(out, "tree = \n"); fprintf(out, "id\tparent\tlchild\trchild\tlsib\trsib\tlikelihood\n"); for (j = 0; j < (int) self->num_nodes; j++) { - fprintf(out, "%d\t%d\t%d\t%d\t%d\t%d\t%f\n", (int) j, self->parent[j], + fprintf(out, "%d\t%d\t%d\t%d\t%d\t%d\t%f\n", (int) j, self->tree[j].parent, self->left_child[j], self->right_child[j], self->left_sib[j], - self->right_sib[j], self->likelihood[j]); + self->right_sib[j], self->likelihood[j].value); } fprintf(out, "likelihood nodes\n"); for (j = 0; j < self->num_likelihood_nodes; j++) { u = self->likelihood_nodes[j]; - fprintf(out, "\t%d -> %f\n", u, self->likelihood[u]); + fprintf(out, "\t%d -> %f\n", u, self->likelihood[u].value); } fprintf(out, "traceback\n"); for (j = 0; j < (int) self->num_sites; j++) { @@ -652,34 +653,26 @@ tsls_hmm_alloc(tsls_hmm_t *self, const tsls_index_t *index, self->traceback = calloc(self->num_sites + 1, sizeof(node_state_list_t)); self->max_likelihood_node = malloc((self->num_sites + 1) * sizeof(tsk_id_t)); - self->parent = tsls_large_alloc((N + 1) * sizeof(*self->parent)); + self->tree = tsls_large_alloc((N + 1) * sizeof(*self->tree)); + self->likelihood = tsls_large_alloc((N + 1) * sizeof(*self->likelihood)); self->left_child = tsls_large_alloc((N + 1) * sizeof(*self->left_child)); self->right_child = tsls_large_alloc((N + 1) * sizeof(*self->right_child)); self->left_sib = tsls_large_alloc((N + 1) * sizeof(*self->left_sib)); self->right_sib = tsls_large_alloc((N + 1) * sizeof(*self->right_sib)); self->num_below = tsls_large_alloc((N + 1) * sizeof(*self->num_below)); self->num_inheriting = tsls_large_alloc((N + 1) * sizeof(*self->num_inheriting)); - self->likelihood = tsls_large_alloc((N + 1) * sizeof(*self->likelihood)); - self->allelic_state = tsls_large_alloc((N + 1) * sizeof(*self->allelic_state)); - self->recombination_required - = tsls_large_alloc((N + 1) * sizeof(*self->recombination_required)); - self->likelihood_cache = tsls_large_alloc((N + 1) * sizeof(*self->likelihood_cache)); - self->recombination_required_cache - = tsls_large_alloc((N + 1) * sizeof(*self->recombination_required_cache)); self->likelihood_nodes = tsls_large_alloc((N + 1) * sizeof(*self->likelihood_nodes)); self->likelihood_nodes_tmp = tsls_large_alloc((N + 1) * sizeof(*self->likelihood_nodes_tmp)); self->values_tmp = tsls_large_alloc((N + 1) * sizeof(*self->values_tmp)); if (self->recombination_rate == NULL || self->mismatch_rate == NULL || self->traceback == NULL - || self->max_likelihood_node == NULL || self->parent == NULL + || self->max_likelihood_node == NULL || self->tree == NULL + || self->likelihood == NULL || self->left_child == NULL || self->right_child == NULL || 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) { + || self->num_inheriting == NULL || self->likelihood_nodes == NULL + || self->likelihood_nodes_tmp == NULL || self->values_tmp == NULL) { ret = TSLS_ERR_NO_MEMORY; goto out; } @@ -699,7 +692,6 @@ tsls_hmm_alloc(tsls_hmm_t *self, const tsls_index_t *index, self->num_sites * sizeof(*self->recombination_rate)); memcpy(self->mismatch_rate, mismatch_rate, self->num_sites * sizeof(*self->mismatch_rate)); - memset(self->allelic_state, 0xff, N * sizeof(*self->allelic_state)); memset(self->max_likelihood_node, 0xff, self->num_sites * sizeof(tsk_id_t)); out: return ret; @@ -710,21 +702,17 @@ tsls_hmm_free(tsls_hmm_t *self) { tsi_safe_free(self->recombination_rate); tsi_safe_free(self->mismatch_rate); - tsls_large_free_array(self->parent, self->num_nodes); + tsls_large_free_array(self->tree, self->num_nodes); + tsls_large_free_array(self->likelihood, self->num_nodes); tsls_large_free_array(self->left_child, self->num_nodes); tsls_large_free_array(self->right_child, self->num_nodes); tsls_large_free_array(self->left_sib, self->num_nodes); tsls_large_free_array(self->right_sib, self->num_nodes); tsls_large_free_array(self->num_below, self->num_nodes); tsls_large_free_array(self->num_inheriting, self->num_nodes); - tsls_large_free_array(self->recombination_required, self->num_nodes); - tsls_large_free_array(self->likelihood, self->num_nodes); - tsls_large_free_array(self->likelihood_cache, self->num_nodes); - tsls_large_free_array(self->recombination_required_cache, self->num_nodes); tsls_large_free_array(self->likelihood_nodes, self->num_nodes); tsls_large_free_array(self->likelihood_nodes_tmp, self->num_nodes); tsls_large_free_array(self->values_tmp, self->num_nodes); - tsls_large_free_array(self->allelic_state, self->num_nodes); tsi_safe_free(self->max_likelihood_node); tsi_safe_free(self->traceback); tsi_blkalloc_free(&self->traceback_allocator); @@ -773,7 +761,7 @@ tsls_hmm_sort_values(tsls_hmm_t *self) int j; for (j = 0; j < self->num_likelihood_nodes; j++) { - self->values_tmp[j] = self->likelihood[self->likelihood_nodes[j]]; + self->values_tmp[j] = self->likelihood[self->likelihood_nodes[j]].value; } qsort(self->values_tmp, (size_t) self->num_likelihood_nodes, sizeof(double), cmp_double); @@ -784,7 +772,7 @@ tsls_hmm_sort_values(tsls_hmm_t *self) static void tsls_hmm_prune(tsls_hmm_t *self) { - double *restrict L = self->likelihood; + tsls_likelihood_t *restrict L = self->likelihood; const tsk_id_t *restrict L_nodes = self->likelihood_nodes; const int num_likelihood_nodes = self->num_likelihood_nodes; double max_L, cutoff; @@ -796,7 +784,7 @@ tsls_hmm_prune(tsls_hmm_t *self) if (self->pruning_mode == TSLS_PRUNE_RELATIVE) { max_L = 0; for (j = 0; j < num_likelihood_nodes; j++) { - max_L = TSK_MAX(max_L, L[L_nodes[j]]); + max_L = TSK_MAX(max_L, L[L_nodes[j]].value); } cutoff = self->pruning_parameter * max_L; } else { @@ -809,8 +797,8 @@ tsls_hmm_prune(tsls_hmm_t *self) cutoff = self->values_tmp[num_likelihood_nodes - n]; } for (j = 0; j < num_likelihood_nodes; j++) { - if (L[L_nodes[j]] < cutoff) { - L[L_nodes[j]] = self->likelihood_threshold; + if (L[L_nodes[j]].value < cutoff) { + L[L_nodes[j]].value = self->likelihood_threshold; } } } @@ -838,7 +826,7 @@ tsls_hmm_count_distinct(tsls_hmm_t *self) static inline void tsls_hmm_add_likelihood(tsls_hmm_t *self, tsk_id_t u, double value) { - self->likelihood[u] = value; + self->likelihood[u].value = value; self->likelihood_nodes[self->num_likelihood_nodes] = u; self->num_likelihood_nodes++; } @@ -848,14 +836,14 @@ tsls_hmm_add_likelihood(tsls_hmm_t *self, tsk_id_t u, double value) static void tsls_hmm_remove_left_nodes(tsls_hmm_t *self) { + tsls_likelihood_t *restrict L = self->likelihood; int j, k; tsk_id_t *restrict L_nodes = self->likelihood_nodes; - const double *restrict L = self->likelihood; k = 0; for (j = 0; j < self->num_likelihood_nodes; j++) { L_nodes[k] = L_nodes[j]; - if (L[L_nodes[j]] != NOT_IN_TREE) { + if (L[L_nodes[j]].value != NOT_IN_TREE) { k++; } } @@ -865,10 +853,10 @@ tsls_hmm_remove_left_nodes(tsls_hmm_t *self) static inline void tsls_hmm_update_num_below(tsls_hmm_t *self, tsk_id_t u, int64_t delta) { - const tsk_id_t *restrict parent = self->parent; + tsls_node_t *restrict tree = self->tree; int64_t *restrict num_below = self->num_below; - for (; u != NULL_NODE; u = parent[u]) { + for (; u != NULL_NODE; u = tree[u].parent) { num_below[u] += delta; } } @@ -881,24 +869,23 @@ tsls_hmm_update_num_below(tsls_hmm_t *self, tsk_id_t u, int64_t delta) static inline double tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u) { - const double *restrict L = self->likelihood; - double *restrict L_cache = self->likelihood_cache; - const tsk_id_t *restrict parent = self->parent; + tsls_likelihood_t *restrict L = self->likelihood; + tsls_node_t *restrict tree = self->tree; double value; tsk_id_t v; v = u; - while (likely(L[v] == NULL_LIKELIHOOD) && likely(L_cache[v] == CACHE_UNSET)) { - v = parent[v]; + while (likely(L[v].value == NULL_LIKELIHOOD) && likely(L[v].cache == CACHE_UNSET)) { + v = tree[v].parent; } - value = L_cache[v]; + value = L[v].cache; if (unlikely(value == CACHE_UNSET)) { - value = L[v]; + value = L[v].value; } v = u; - while (likely(L[v] == NULL_LIKELIHOOD) && likely(L_cache[v] == CACHE_UNSET)) { - L_cache[v] = value; - v = parent[v]; + while (likely(L[v].value == NULL_LIKELIHOOD) && likely(L[v].cache == CACHE_UNSET)) { + L[v].cache = value; + v = tree[v].parent; } return value; } @@ -909,12 +896,12 @@ tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u) static inline void tsls_hmm_reset_likelihood_cache(tsls_hmm_t *self, tsk_id_t u) { - double *restrict L_cache = self->likelihood_cache; - const tsk_id_t *restrict parent = self->parent; + tsls_likelihood_t *restrict L = self->likelihood; + tsls_node_t *restrict tree = self->tree; - while (L_cache[u] != CACHE_UNSET) { - L_cache[u] = CACHE_UNSET; - u = parent[u]; + while (L[u].cache != CACHE_UNSET) { + L[u].cache = CACHE_UNSET; + u = tree[u].parent; } } @@ -927,23 +914,23 @@ static inline void tsls_hmm_remove_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, bool update_likelihoods, tsk_id_t *restrict isolated, int *num_isolated) { - double *restrict L = self->likelihood; + tsls_likelihood_t *restrict L = self->likelihood; if (self->track_num_below) { tsls_hmm_update_num_below(self, p, -self->num_below[c]); } - if (update_likelihoods && L[c] == NULL_LIKELIHOOD) { + if (update_likelihoods && L[c].value == NULL_LIKELIHOOD) { tsls_hmm_add_likelihood(self, c, tsls_hmm_get_inherited_likelihood(self, p)); } - remove_edge(p, c, self->parent, self->left_child, self->right_child, + remove_edge(p, c, self->tree, self->left_child, self->right_child, self->left_sib, self->right_sib); - self->allelic_state[c] = ROOT_ALLELE; + self->tree[c].allelic_state = ROOT_ALLELE; /* c is now a root, and so has a likelihood. Its cached value (if any, * from an earlier removal at this position) is the same as its * likelihood, and clearing it means that no root has a cached value, * so walks over the cache always stop below a root (see * tsls_hmm_reset_likelihood_cache). */ - self->likelihood_cache[c] = CACHE_UNSET; + self->likelihood[c].cache = CACHE_UNSET; if (update_likelihoods) { /* Each node can only become isolated once when removing edges, so * the isolated list has at most num_nodes entries. */ @@ -953,7 +940,7 @@ tsls_hmm_remove_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, } if (!tsls_hmm_in_tree(self, p)) { /* p was a root, and so has a likelihood */ - assert(L[p] >= 0); + assert(L[p].value >= 0); isolated[*num_isolated] = p; (*num_isolated)++; } @@ -965,21 +952,21 @@ tsls_hmm_remove_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, static inline void tsls_hmm_insert_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, bool update_likelihoods) { - double *restrict L = self->likelihood; + tsls_likelihood_t *restrict L = self->likelihood; if (update_likelihoods) { - if (L[p] == NOT_IN_TREE) { + if (L[p].value == NOT_IN_TREE) { tsls_hmm_add_likelihood(self, p, 0); } - if (L[c] == NOT_IN_TREE) { + if (L[c].value == NOT_IN_TREE) { tsls_hmm_add_likelihood(self, c, 0); } /* Otherwise the child was a root and keeps its likelihood */ - assert(L[c] >= 0); + assert(L[c].value >= 0); } - insert_edge(p, c, self->parent, self->left_child, self->right_child, + insert_edge(p, c, self->tree, self->left_child, self->right_child, self->left_sib, self->right_sib); - self->allelic_state[c] = TSK_NULL; + self->tree[c].allelic_state = TSK_NULL; if (self->track_num_below) { tsls_hmm_update_num_below(self, p, self->num_below[c]); } @@ -990,11 +977,11 @@ tsls_hmm_insert_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, bool update_likel static void tsls_hmm_update_tree(tsls_hmm_t *self, const tsk_tree_position_t *tp) { + tsls_likelihood_t *restrict L = self->likelihood; const tsls_edge_t *restrict out = tsls_index_get_edges(self->index, tp->out.order); const tsls_edge_t *restrict in = tsls_index_get_edges(self->index, tp->in.order); const int direction = tp->direction; tsk_id_t *restrict isolated = self->likelihood_nodes_tmp; - double *restrict L = self->likelihood; tsk_id_t j, u; int k, num_left, num_isolated = 0; @@ -1016,7 +1003,7 @@ tsls_hmm_update_tree(tsls_hmm_t *self, const tsk_tree_position_t *tp) for (k = 0; k < num_isolated; k++) { u = isolated[k]; if (!tsls_hmm_in_tree(self, u)) { - L[u] = NOT_IN_TREE; + L[u].value = NOT_IN_TREE; num_left++; } } @@ -1038,16 +1025,17 @@ tsls_hmm_init_sweep( tsk_id_t u, j; size_t k; - memset(self->parent, 0xff, N * sizeof(*self->parent)); memset(self->left_child, 0xff, N * sizeof(*self->left_child)); memset(self->right_child, 0xff, N * sizeof(*self->right_child)); memset(self->left_sib, 0xff, N * sizeof(*self->left_sib)); memset(self->right_sib, 0xff, N * sizeof(*self->right_sib)); - memset(self->recombination_required, 0xff, N * sizeof(*self->recombination_required)); - /* Every node starts as a root (ROOT_ALLELE is 0xfe as an int8) */ - memset(self->allelic_state, (uint8_t) ROOT_ALLELE, N * sizeof(*self->allelic_state)); for (k = 0; k < N; k++) { - self->likelihood_cache[k] = CACHE_UNSET; + /* Every node starts as a root. The likelihoods are set below. */ + self->likelihood[k].cache = CACHE_UNSET; + self->tree[k].parent = NULL_NODE; + self->tree[k].allelic_state = ROOT_ALLELE; + self->tree[k].recombination_required = -1; + self->tree[k].recombination_required_cache = -1; self->num_below[k] = 1; } ret = seek_tree_position(tp, position); @@ -1062,12 +1050,12 @@ tsls_hmm_init_sweep( } self->num_likelihood_nodes = 0; for (u = 0; u < (tsk_id_t) N; u++) { - if (self->parent[u] != NULL_NODE) { - self->likelihood[u] = NULL_LIKELIHOOD; + if (self->tree[u].parent != NULL_NODE) { + self->likelihood[u].value = NULL_LIKELIHOOD; } else if (self->left_child[u] != NULL_NODE) { tsls_hmm_add_likelihood(self, u, initial_value); } else { - self->likelihood[u] = NOT_IN_TREE; + self->likelihood[u].value = NOT_IN_TREE; } } out: @@ -1084,7 +1072,7 @@ tsls_hmm_set_allelic_state(tsls_hmm_t *self, const tsk_id_t site) tsk_size_t j; for (j = index->site_mutations_offset[site]; j < stop; j++) { - self->allelic_state[node[j]] = derived_state[j]; + self->tree[node[j]].allelic_state = derived_state[j]; } } @@ -1101,7 +1089,7 @@ tsls_hmm_unset_allelic_state(tsls_hmm_t *self, const tsk_id_t site) for (j = index->site_mutations_offset[site]; j < stop; j++) { u = node[j]; - self->allelic_state[u] = self->parent[u] == NULL_NODE ? ROOT_ALLELE : TSK_NULL; + self->tree[u].allelic_state = self->tree[u].parent == NULL_NODE ? ROOT_ALLELE : TSK_NULL; } } @@ -1116,7 +1104,7 @@ tsls_hmm_clear_allelic_state(tsls_hmm_t *self, const tsk_id_t site) tsk_size_t j; for (j = index->site_mutations_offset[site]; j < stop; j++) { - self->allelic_state[node[j]] = TSK_NULL; + self->tree[node[j]].allelic_state = TSK_NULL; } } @@ -1125,8 +1113,8 @@ tsls_hmm_clear_allelic_state(tsls_hmm_t *self, const tsk_id_t site) static void tsls_hmm_insert_mutation_likelihoods(tsls_hmm_t *self, const tsk_id_t site) { - const double *restrict L = self->likelihood; - const tsk_id_t *restrict parent = self->parent; + tsls_likelihood_t *restrict L = self->likelihood; + tsls_node_t *restrict tree = self->tree; const tsls_index_t *index = self->index; const tsk_size_t stop = index->site_mutations_offset[site + 1]; tsk_size_t j; @@ -1134,12 +1122,12 @@ tsls_hmm_insert_mutation_likelihoods(tsls_hmm_t *self, const tsk_id_t site) for (j = index->site_mutations_offset[site]; j < stop; j++) { node = index->mutation_node[j]; - if (L[node] == NULL_LIKELIHOOD) { + if (L[node].value == NULL_LIKELIHOOD) { u = node; - while (L[u] == NULL_LIKELIHOOD) { - u = parent[u]; + while (L[u].value == NULL_LIKELIHOOD) { + u = tree[u].parent; } - tsls_hmm_add_likelihood(self, node, L[u]); + tsls_hmm_add_likelihood(self, node, L[u].value); } } } @@ -1151,31 +1139,28 @@ static ALWAYS_INLINE double tsls_hmm_get_inherited_likelihood_and_flag( tsls_hmm_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; + tsls_likelihood_t *restrict L = self->likelihood; + tsls_node_t *restrict tree = self->tree; 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]; + while (likely(L[v].value == NULL_LIKELIHOOD) && likely(L[v].cache == CACHE_UNSET)) { + v = tree[v].parent; } - if (L_cache[v] == CACHE_UNSET) { - value = L[v]; - r = R[v]; + if (L[v].cache == CACHE_UNSET) { + value = L[v].value; + r = tree[v].recombination_required; } else { - value = L_cache[v]; - r = R_cache[v]; + value = L[v].cache; + r = tree[v].recombination_required_cache; } 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]; + while (likely(L[v].value == NULL_LIKELIHOOD) && likely(L[v].cache == CACHE_UNSET)) { + L[v].cache = value; + tree[v].recombination_required_cache = r; + v = tree[v].parent; } *flag = r; return value; @@ -1190,12 +1175,11 @@ static ALWAYS_INLINE void tsls_hmm_coalesce_likelihoods( tsls_hmm_t *self, bool check_recombination_required) { - const int8_t *restrict R = self->recombination_required; + tsls_likelihood_t *restrict L = self->likelihood; + tsls_node_t *restrict tree = self->tree; 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; tsk_id_t *restrict cached_paths = self->likelihood_nodes_tmp; const int old_num_likelihood_nodes = self->num_likelihood_nodes; tsk_id_t *restrict L_nodes = self->likelihood_nodes; @@ -1205,7 +1189,7 @@ tsls_hmm_coalesce_likelihoods( num_likelihood_nodes = 0; for (j = 0; j < old_num_likelihood_nodes; j++) { u = L_nodes[j]; - p = parent[u]; + p = tree[u].parent; if (p != NULL_NODE) { cached_paths[num_cached_paths] = p; num_cached_paths++; @@ -1214,17 +1198,17 @@ tsls_hmm_coalesce_likelihoods( * cached values below it stay valid. */ L_p = tsls_hmm_get_inherited_likelihood_and_flag( self, p, &R_p); - if (L[u] == L_p && R[u] == R_p) { - L[u] = NULL_LIKELIHOOD; + if (L[u].value == L_p && tree[u].recombination_required == R_p) { + L[u].value = NULL_LIKELIHOOD; } } else { L_p = tsls_hmm_get_inherited_likelihood(self, p); - if (L[u] == L_p) { - L[u] = NULL_LIKELIHOOD; + if (L[u].value == L_p) { + L[u].value = NULL_LIKELIHOOD; } } } - if (L[u] >= 0) { + if (L[u].value >= 0) { L_nodes[num_likelihood_nodes] = L_nodes[j]; num_likelihood_nodes++; } @@ -1270,8 +1254,8 @@ tsls_hmm_get_n(const tsls_hmm_t *self) static void tsls_hmm_compute_num_inheriting(tsls_hmm_t *self) { - const double *restrict L = self->likelihood; - const tsk_id_t *restrict parent = self->parent; + tsls_likelihood_t *restrict L = self->likelihood; + tsls_node_t *restrict tree = self->tree; const tsk_id_t *restrict L_nodes = self->likelihood_nodes; const int64_t *restrict num_below = self->num_below; int64_t *restrict num_inheriting = self->num_inheriting; @@ -1284,10 +1268,10 @@ tsls_hmm_compute_num_inheriting(tsls_hmm_t *self) } for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; - w = parent[u]; + w = tree[u].parent; if (w != NULL_NODE) { - while (L[w] == NULL_LIKELIHOOD) { - w = parent[w]; + while (L[w].value == NULL_LIKELIHOOD) { + w = tree[w].parent; } num_inheriting[w] -= num_below[u]; } @@ -1298,7 +1282,7 @@ tsls_hmm_compute_num_inheriting(tsls_hmm_t *self) static double tsls_hmm_total_likelihood(const tsls_hmm_t *self) { - const double *restrict L = self->likelihood; + tsls_likelihood_t *restrict L = self->likelihood; const tsk_id_t *restrict L_nodes = self->likelihood_nodes; const int64_t *restrict num_inheriting = self->num_inheriting; double total = 0; @@ -1308,7 +1292,7 @@ tsls_hmm_total_likelihood(const tsls_hmm_t *self) for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; assert(num_inheriting[u] > 0); - total += (double) num_inheriting[u] * L[u]; + total += (double) num_inheriting[u] * L[u].value; } return total; } @@ -1325,7 +1309,7 @@ tsls_hmm_store_traceback(tsls_hmm_t *self, const tsk_id_t site_id) node_state_list_t *restrict list; node_state_list_t *restrict T = self->traceback; const tsk_id_t *restrict nodes = self->likelihood_nodes; - const int8_t *restrict R = self->recombination_required; + tsls_node_t *restrict tree = self->tree; const int num_likelihood_nodes = self->num_likelihood_nodes; bool match; @@ -1339,7 +1323,7 @@ tsls_hmm_store_traceback(tsls_hmm_t *self, const tsk_id_t site_id) list_R = list->recombination_required; match = true; for (j = 0; j < num_likelihood_nodes; j++) { - if (list_node[j] != nodes[j] || list_R[j] != R[nodes[j]]) { + if (list_node[j] != nodes[j] || list_R[j] != tree[nodes[j]].recombination_required) { match = false; break; } @@ -1366,7 +1350,7 @@ tsls_hmm_store_traceback(tsls_hmm_t *self, const tsk_id_t site_id) for (j = 0; j < num_likelihood_nodes; j++) { u = nodes[j]; list_node[j] = u; - list_R[j] = R[u]; + list_R[j] = tree[u].recombination_required; } } self->total_traceback_size += (size_t) num_likelihood_nodes; @@ -1378,13 +1362,11 @@ static int WARN_UNUSED tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site, const allele_t state, compressed_matrix_t *output) { + tsls_likelihood_t *restrict L = self->likelihood; int ret = 0; const int num_likelihood_nodes = self->num_likelihood_nodes; 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; - double *restrict L = self->likelihood; - int8_t *restrict recombination_required = self->recombination_required; + tsls_node_t *restrict tree = self->tree; const double rho = self->recombination_rate[site]; const double mu = self->mismatch_rate[site]; const double n = tsls_hmm_get_n(self); @@ -1400,21 +1382,21 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site, max_L_node = NULL_NODE; for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; - p_no_recomb = L[u] * (1 - rho + rho / n); + p_no_recomb = L[u].value * (1 - rho + rho / n); p_recomb = rho / n; - recombination_required[u] = false; + tree[u].recombination_required = false; if (p_no_recomb > p_recomb) { p_t = p_no_recomb; } else { p_t = p_recomb; - recombination_required[u] = true; + tree[u].recombination_required = true; } p_e = emission_probability( - get_marked_allelic_state(allelic_state, parent, u, ancestral_state), state, + get_marked_allelic_state(tree, u, ancestral_state), state, mu, num_alleles); - L[u] = p_t * p_e; - if (L[u] > max_L) { - max_L = L[u]; + L[u].value = p_t * p_e; + if (L[u].value > max_L) { + max_L = L[u].value; max_L_node = u; } } @@ -1435,7 +1417,7 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site, self->max_likelihood_node[site] = max_L_node; for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; - L[u] = TSK_MAX(L[u] / max_L, self->likelihood_threshold); + L[u].value = TSK_MAX(L[u].value / max_L, self->likelihood_threshold); } tsls_hmm_prune(self); tsls_hmm_count_distinct(self); @@ -1448,8 +1430,8 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site, goto out; } if (output != NULL) { - ret = compressed_matrix_store_site(output, site, max_L, - self->num_likelihood_nodes, L_nodes, L, recombination_required); + ret = compressed_matrix_store_site( + output, site, max_L, self->num_likelihood_nodes, L_nodes, L, self->tree, true); if (ret != 0) { goto out; } @@ -1466,11 +1448,10 @@ static int WARN_UNUSED tsls_hmm_process_site_forward(tsls_hmm_t *self, const tsk_id_t site, const allele_t state, double *previous_total, compressed_matrix_t *output) { + tsls_likelihood_t *restrict L = self->likelihood; int ret = 0; 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; - double *restrict L = self->likelihood; + tsls_node_t *restrict tree = self->tree; const double rho = self->recombination_rate[site]; const double mu = self->mismatch_rate[site]; const double n = tsls_hmm_get_n(self); @@ -1490,9 +1471,9 @@ tsls_hmm_process_site_forward(tsls_hmm_t *self, const tsk_id_t site, for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; p_e = emission_probability( - get_marked_allelic_state(allelic_state, parent, u, ancestral_state), state, + get_marked_allelic_state(tree, u, ancestral_state), state, mu, num_alleles); - L[u] = (L[u] * (1 - rho) + rho / n * total) * p_e; + L[u].value = (L[u].value * (1 - rho) + rho / n * total) * p_e; } tsls_hmm_unset_allelic_state(self, site); normalisation_factor = tsls_hmm_total_likelihood(self); @@ -1502,12 +1483,12 @@ tsls_hmm_process_site_forward(tsls_hmm_t *self, const tsk_id_t site, } for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; - L[u] = TSK_MAX(L[u] / normalisation_factor, self->likelihood_threshold); + L[u].value = TSK_MAX(L[u].value / normalisation_factor, self->likelihood_threshold); } tsls_hmm_prune(self); *previous_total = tsls_hmm_total_likelihood(self); ret = compressed_matrix_store_site(output, site, normalisation_factor, - self->num_likelihood_nodes, L_nodes, L, NULL); + self->num_likelihood_nodes, L_nodes, L, self->tree, false); if (ret != 0) { goto out; } @@ -1532,11 +1513,10 @@ tsls_hmm_process_site_backward(tsls_hmm_t *self, const tsk_id_t site, const allele_t state, const double *forward_normalisation_factor, backward_transition_t *transition, compressed_matrix_t *output) { + tsls_likelihood_t *restrict L = self->likelihood; int ret = 0; 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; - double *restrict L = self->likelihood; + tsls_node_t *restrict tree = self->tree; const double mu = self->mismatch_rate[site]; const double n = tsls_hmm_get_n(self); const double num_alleles = (double) self->index->num_alleles[site]; @@ -1548,15 +1528,15 @@ tsls_hmm_process_site_backward(tsls_hmm_t *self, const tsk_id_t site, if (transition->pending) { for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; - x = (1 - transition->rho) * L[u] + transition->rho * transition->total / n; - L[u] = TSK_MAX( + x = (1 - transition->rho) * L[u].value + transition->rho * transition->total / n; + L[u].value = TSK_MAX( x / transition->normalisation_factor, self->likelihood_threshold); } tsls_hmm_prune(self); } tsls_hmm_insert_mutation_likelihoods(self, site); ret = compressed_matrix_store_site(output, site, forward_normalisation_factor[site], - self->num_likelihood_nodes, L_nodes, L, NULL); + self->num_likelihood_nodes, L_nodes, L, self->tree, false); if (ret != 0) { goto out; } @@ -1564,8 +1544,8 @@ tsls_hmm_process_site_backward(tsls_hmm_t *self, const tsk_id_t site, tsls_hmm_set_allelic_state(self, site); for (j = 0; j < self->num_likelihood_nodes; j++) { u = L_nodes[j]; - L[u] *= emission_probability( - get_marked_allelic_state(allelic_state, parent, u, ancestral_state), state, + L[u].value *= emission_probability( + get_marked_allelic_state(tree, u, ancestral_state), state, mu, num_alleles); } tsls_hmm_unset_allelic_state(self, site); @@ -1642,15 +1622,15 @@ tsls_hmm_run(tsls_hmm_t *self, int mode, tsk_id_t start, tsk_id_t end, */ static inline void tsls_hmm_set_recombination_required( - tsls_hmm_t *self, tsk_id_t site, int8_t *restrict recombination_required) + tsls_hmm_t *self, tsk_id_t site, tsls_node_t *restrict tree) { int j; - const int8_t *restrict R = self->traceback[site].recombination_required; - const tsk_id_t *restrict node = self->traceback[site].node; + const int8_t *restrict list_R = self->traceback[site].recombination_required; + const tsk_id_t *restrict list_node = self->traceback[site].node; const int size = self->traceback[site].size; for (j = 0; j < size; j++) { - recombination_required[node[j]] = R[j]; + tree[list_node[j]].recombination_required = list_R[j]; } } @@ -1658,14 +1638,14 @@ tsls_hmm_set_recombination_required( */ static inline void tsls_hmm_unset_recombination_required( - tsls_hmm_t *self, tsk_id_t site, int8_t *restrict recombination_required) + tsls_hmm_t *self, tsk_id_t site, tsls_node_t *restrict tree) { int j; - const tsk_id_t *restrict node = self->traceback[site].node; + const tsk_id_t *restrict list_node = self->traceback[site].node; const int size = self->traceback[site].size; for (j = 0; j < size; j++) { - recombination_required[node[j]] = -1; + tree[list_node[j]].recombination_required = -1; } } @@ -1681,9 +1661,7 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, tsk_id_t site, j, u, max_likelihood_node; double x; bool valid; - tsk_id_t *restrict parent = self->parent; - allele_t *restrict allelic_state = self->allelic_state; - int8_t *restrict recombination_required = self->recombination_required; + tsls_node_t *restrict tree = self->tree; const double *restrict sites_position = self->sites_position; /* Every root of the traceback tree has the virtual root as its parent, * whose allelic state is the site's ancestral state, so the walk to find @@ -1708,12 +1686,11 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, /* Build the tree at the last site, and then go through the trees in * reverse, running the traceback at each site. */ for (u = 0; u < virtual_root; u++) { - parent[u] = virtual_root; + tree[u].parent = virtual_root; + tree[u].allelic_state = TSK_NULL; + tree[u].recombination_required = -1; } - parent[virtual_root] = NULL_NODE; - memset( - recombination_required, 0xff, self->num_nodes * sizeof(*recombination_required)); - memset(allelic_state, 0xff, self->num_nodes * sizeof(*allelic_state)); + tree[virtual_root].parent = NULL_NODE; x = sites_position[end - 1]; ret = seek_tree_position(&tp, x); if (ret != 0) { @@ -1722,7 +1699,7 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, in = tsls_index_get_edges(self->index, tp.in.order); for (j = tp.in.start; j != tp.in.stop; j += tp.direction) { if (seek_edge_covers(&tp, &in[j], x)) { - parent[in[j].child] = in[j].parent; + tree[in[j].child].parent = in[j].parent; } } for (site = end - 1; site >= start; site--) { @@ -1733,33 +1710,33 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, out = tsls_index_get_edges(self->index, tp.out.order); in = tsls_index_get_edges(self->index, tp.in.order); for (j = tp.out.start; j != tp.out.stop; j--) { - parent[out[j].child] = virtual_root; + tree[out[j].child].parent = virtual_root; } for (j = tp.in.start; j != tp.in.stop; j--) { - parent[in[j].child] = in[j].parent; + tree[in[j].child].parent = in[j].parent; } } - allelic_state[virtual_root] = self->index->ancestral_state[site]; + tree[virtual_root].allelic_state = self->index->ancestral_state[site]; tsls_hmm_set_allelic_state(self, site); u = path_parent[path_length]; - while (allelic_state[u] == TSK_NULL) { - u = parent[u]; + while (tree[u].allelic_state == TSK_NULL) { + u = tree[u].parent; } - match[site] = allelic_state[u]; + match[site] = tree[u].allelic_state; tsls_hmm_clear_allelic_state(self, site); u = path_parent[path_length]; /* Mark the traceback nodes on the tree */ - tsls_hmm_set_recombination_required(self, site, recombination_required); + tsls_hmm_set_recombination_required(self, site, tree); /* Traverse up the tree from the current node. The first marked * node that we meet tells us whether we need to recombine. Every * root has a likelihood, and so is marked. */ - while (recombination_required[u] == -1) { - u = parent[u]; + while (tree[u].recombination_required == -1) { + u = tree[u].parent; assert(u != virtual_root); } - if (recombination_required[u] && site > start) { + if (tree[u].recombination_required && site > start) { max_likelihood_node = self->max_likelihood_node[site - 1]; assert(max_likelihood_node != NULL_NODE); path_left[path_length] = sites_position[site]; @@ -1769,7 +1746,7 @@ tsls_hmm_run_traceback(tsls_hmm_t *self, tsk_id_t start, tsk_id_t end, path_parent[path_length] = max_likelihood_node; } /* Unset the values in the tree for the next site. */ - tsls_hmm_unset_recombination_required(self, site, recombination_required); + tsls_hmm_unset_recombination_required(self, site, tree); } path_left[path_length] = start_pos; path_length++; @@ -1842,7 +1819,6 @@ tsls_hmm_find_path_matrix(tsls_hmm_t *self, tsk_id_t start, memset(self->max_likelihood_node + start, 0xff, (size_t) (end - start) * sizeof(*self->max_likelihood_node)); } - memset(self->allelic_state, 0xff, self->num_nodes * sizeof(*self->allelic_state)); return ret; } @@ -1870,7 +1846,6 @@ tsls_hmm_forward( haplotype, NULL, output); } out: - memset(self->allelic_state, 0xff, self->num_nodes * sizeof(*self->allelic_state)); return ret; } @@ -1896,7 +1871,6 @@ tsls_hmm_backward(tsls_hmm_t *self, const allele_t *haplotype, haplotype, forward_normalisation_factor, output); } out: - memset(self->allelic_state, 0xff, self->num_nodes * sizeof(*self->allelic_state)); return ret; } diff --git a/tsls/lib/tsls.h b/tsls/lib/tsls.h index 38ad3c4..4ecde61 100644 --- a/tsls/lib/tsls.h +++ b/tsls/lib/tsls.h @@ -130,6 +130,27 @@ typedef struct { tsls_edge_t *removal_edges; } tsls_index_t; +/* The state of a node that the walks up the tree read, kept together (in + * 8 bytes) so that each step of a walk reads one record. */ +typedef struct { + tsk_id_t parent; + /* The allelic state at the current site, if the node has a mutation + * there (or is a root, during the sweeps) */ + allele_t allelic_state; + int8_t recombination_required; + /* The recombination_required flag inherited by the node, valid where + * the likelihood cache is set. Only filled when coalescing Viterbi. */ + int8_t recombination_required_cache; +} tsls_node_t; + +/* A node's likelihood, with the cache used when coalescing. */ +typedef struct { + /* The node's likelihood, or NULL_LIKELIHOOD or NOT_IN_TREE */ + double value; + /* The likelihood the node inherits, cached during coalescing */ + double cache; +} tsls_likelihood_t; + /* The Li and Stephens HMM on a tree sequence. The index is borrowed, and * must not be changed or freed while the HMM is in use. */ typedef struct { @@ -146,8 +167,10 @@ typedef struct { double pruning_parameter; double *recombination_rate; double *mismatch_rate; - /* The quintuply linked tree */ - tsk_id_t *parent; + /* The per-node state, with the parent links of the quintuply linked + * tree, which are completed by the arrays below */ + tsls_node_t *tree; + tsls_likelihood_t *likelihood; tsk_id_t *left_child; tsk_id_t *right_child; tsk_id_t *left_sib; @@ -158,17 +181,9 @@ typedef struct { bool track_num_below; /* The number of nodes in the tree that inherit each likelihood value. */ 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. */ tsk_id_t *max_likelihood_node; - /* Used during traceback to map nodes where recombination is required. */ - int8_t *recombination_required; tsk_id_t *likelihood_nodes_tmp; tsk_id_t *likelihood_nodes; node_state_list_t *traceback; From f0455baabd342834de8225543c3c6c4512c4ca56 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 16:30:17 +0100 Subject: [PATCH 24/25] Update the README for the two node records The current result is now 0.78-0.90 of tsinfer-hmm's wall time on every workload, with a column for the node records in the progression table, and the layouts section notes that the two-record layout only helped once huge pages were used. --- benchmarking/README.md | 68 +++++++++++++++++++++++++----------------- 1 file changed, 40 insertions(+), 28 deletions(-) diff --git a/benchmarking/README.md b/benchmarking/README.md index d60c226..248c1b3 100644 --- a/benchmarking/README.md +++ b/benchmarking/README.md @@ -502,34 +502,36 @@ uv run python paired.py #### Current result -Two full runs of the current build (`9806f55`), tsls / tsinfer-hmm: +Two full runs of the current build (`832fc5f`), tsls / tsinfer-hmm: | workload | queries | round 1 | round 2 | |---|---|---|---| -| ancestors, group 84 | 50 | 0.72 | 0.75 | -| ancestors, group 1000 | 50 | 0.76 | 0.74 | -| ancestors, group 1483 | 100 | 0.91 | 0.80 | -| ancestors, group 1484 | 50 | 0.96 | 0.79 | -| ancestors, group 1485 | 100 | 1.00 | 1.00 | -| samples | 6 | 0.87 | 0.81 | - -tsls is now faster than tsinfer-hmm on samples (by 13-19%) and on most -ancestor groups, and level on group 1485, with the same paths and -traceback sizes. +| ancestors, group 84 | 50 | 0.80 | 0.81 | +| ancestors, group 1000 | 50 | 0.82 | 0.82 | +| ancestors, group 1483 | 100 | 0.78 | 0.78 | +| ancestors, group 1484 | 50 | 0.84 | 0.82 | +| ancestors, group 1485 | 100 | 0.90 | 0.87 | +| samples | 6 | 0.79 | 0.82 | + +tsls is now 10-22% faster than tsinfer-hmm on every workload, with the +same paths and traceback sizes. It executes more instructions than +tsinfer-hmm on samples again (1.13 times, on the profiling sample), but +fewer on all the ancestor groups. #### How it got there -| workload | after the refactor | + shared index | + huge pages | -|---|---|---|---| -| ancestors, group 84 | 1.43 | 0.73-0.77 | 0.72-0.75 | -| ancestors, group 1000 | 1.19 | 0.75-0.76 | 0.74-0.76 | -| ancestors, group 1483 | 1.19 | 0.88-0.92 | 0.80-0.91 | -| ancestors, group 1484 | 1.23 | 0.99-1.00 | 0.79-0.96 | -| ancestors, group 1485 | 1.25 | 0.97-1.06 | 1.00 | -| samples | 1.14 | 0.99-1.17 | 0.81-0.87 | +| workload | after the refactor | + shared index | + huge pages | + two node records | +|---|---|---|---|---| +| ancestors, group 84 | 1.43 | 0.73-0.77 | 0.72-0.76 | 0.80-0.81 | +| ancestors, group 1000 | 1.19 | 0.75-0.76 | 0.74-0.76 | 0.82 | +| ancestors, group 1483 | 1.19 | 0.88-0.92 | 0.80-0.96 | 0.78 | +| ancestors, group 1484 | 1.23 | 0.99-1.00 | 0.79-0.96 | 0.82-0.84 | +| ancestors, group 1485 | 1.25 | 0.97-1.06 | 0.95-1.00 | 0.87-0.90 | +| samples | 1.14 | 0.99-1.17 | 0.81-1.11 | 0.79-0.82 | -(Ranges are over four runs of the shared index build and two of the huge -pages build, run alternately with other builds.) +(Ranges are over four runs each of the shared index and huge pages +builds, and two of the current build, run alternately with other +builds.) - **After the refactor**, tsls was slower despite executing fewer instructions. Building a `tsls.Hmm` took 4-8 ms against 1.5 ms for @@ -553,6 +555,15 @@ pages build, run alternately with other builds.) directly: glibc reuses memory already faulted in with small pages, which `madvise` doesn't change, so `posix_memalign` + `madvise` gave no huge pages at all. +- **Two node records** (`832fc5f`): the per-node fields the walks read are + kept in an 8-byte record {parent, allelic state, recombination flag and + its cache} and a 16-byte record {likelihood, its cache}, so the allele + walk reads one record per step and the likelihood cache walks two, + instead of two and four arrays. This helps the late groups and samples + and costs a little on the early groups, whose queries take only 4-8 ms. + It also costs instructions (the profiling sample goes from 1.52G to + 1.76G; GCC forms a pointer to the record at each step of the walk), but + each step is still one memory access. #### Measuring on this machine @@ -575,7 +586,7 @@ pages build, run alternately with other builds.) writes (the cache clear in `remove_edge`, and keeping roots marked through the sweep) made no measurable difference to the samples. -#### Node layouts that didn't help +#### Node layouts A `perf` profile of the shared index build put the samples' remaining time in the likelihood cache walks, which read several per-node arrays at @@ -583,9 +594,10 @@ each step. Two layouts that put these fields together were tried, each in two full runs alternating with the then-current build: - One 24-byte record per node (likelihood, its cache, parent, allelic - state, recombination flag and its cache): worse everywhere (samples - 1.52 and 1.52; ancestors up to 1.64). Indexing 24-byte records also - costs an address computation at every step of the walks: instructions - on the profiling sample rose from 1.52G to 1.82G. -- An 8-byte record {parent, allelic state, flags} and a 16-byte record - {likelihood, its cache}: no consistent difference from separate arrays. + state, recombination flag and its cache): worse everywhere, before huge + pages (samples 1.52 and 1.52; ancestors up to 1.64). Indexing 24-byte + records also costs an address computation at every step of the walks: + instructions on the profiling sample rose from 1.52G to 1.82G. +- The two records now used: no consistent difference from separate arrays + before huge pages, but faster on the late groups and samples with them + (above). From 401daa74dc4e03be35fe6507af583c00c7305c0f Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 23 Oct 2026 16:30:28 +0100 Subject: [PATCH 25/25] Correct the README on run-to-run variation with huge pages Later runs of the huge pages build varied as much as before (samples 14.2-18.0 s over four runs); only the two-record build's two runs were close (14.1 and 14.5 s). --- benchmarking/README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/benchmarking/README.md b/benchmarking/README.md index 248c1b3..bdf9656 100644 --- a/benchmarking/README.md +++ b/benchmarking/README.md @@ -569,8 +569,9 @@ builds.) - Single runs mislead. tsls's time for the same build varied by up to 20% between processes (6 samples: 16.4 to 19.9 s) while tsinfer-hmm's was - steady, which fits physical page placement varying between runs; with - huge pages tsls's times are steadier (14.2 and 14.7 s). Another session + steady, which fits physical page placement varying between runs. With + huge pages alone it still varied (14.2 to 18.0 s over four runs); with + the two node records, the two runs gave 14.1 and 14.5 s. Another session on the machine may also compete for memory bandwidth at times. Builds were therefore compared with full runs repeated in alternation (build A, build B, A, B), and a single `--quick` run once showed 1.36 for group