diff --git a/benchmarking/README.md b/benchmarking/README.md index 5893584..74abc70 100644 --- a/benchmarking/README.md +++ b/benchmarking/README.md @@ -7,8 +7,10 @@ Benchmarks for the Li and Stephens HMM engines in tsls: Each engine runs Viterbi matching over a set of query haplotypes with known ground truth, and the results are recorded so the engines can be -characterised across a range of queries. The first target is the sc2ts -SARS-CoV-2 ARG. Benchmarks for ARGs inferred by tsinfer will follow. +characterised across a range of queries. There are two targets: the +sc2ts SARS-CoV-2 ARG, and the ancestor and sample matching phases of a +tsinfer 1.0 inference of 1000 Genomes chr20p (see +[below](#tsinfer-10-1000-genomes-chr20p)). ## Setup @@ -228,3 +230,143 @@ are the 110 queries from the sub-ARG section (10 nodes, k = 0, 1, 2, 5, large ARGs. - For clean timings, run sub-ARGs up to about 300k samples one at a time. Restrict tskit to the smallest sub-ARGs with `--limit`. + +## tsinfer 1.0: 1000 Genomes chr20p + +These benchmarks replay the HMM calls of a real tsinfer 1.0 inference, so +that the engines can be compared with tsinfer on the workloads it +actually runs. The paths tsinfer found are the ground truth. + +### Data + +A colleague ran tsinfer at commit `cb5017e` on the 1000 Genomes high +coverage chr20p (3,202 individuals, 580,909 inference sites), with the +parameters in `full_inference.toml`: `rho = 0.01` and `mu = 1e-20` for +both ancestors and samples, no path compression, and tsinfer's defaults +of `weight_by_n=True` and no likelihood threshold. The data are copied +from `tmp/for_jerome` to `data/tgp_chr20p/`, leaving out macOS metadata: + +``` +rsync -a --exclude '._*' --exclude .DS_Store ../tmp/for_jerome/ data/tgp_chr20p/ +``` + +- `tgp_chr20p_inferred_ancestors.zarr`: the 432,851 ancestors. +- `tgp_chr20p_inferred_ancestors.trees`: the ancestors tree sequence + (433k nodes, 4.2M edges). +- `tgp_chr20p_samples_masked.zarr`: the sample genotypes. +- `tgp_chr20p_inferred_samples_raw.trees`: the result of matching the + 6,404 sample haplotypes, before post-processing. +- `partial_ancestor_match/`: the tree sequence before group 84 + (`before.trees`), the group's ancestors (`group.zarr`, `group.json`) and + the result (`after.trees`). + +See `data/tgp_chr20p/README.md` for details. + +### Workflow + +#### 1. Compute the match groups + +The tree sequences do not record which nodes were matched in which +group, so we recompute the groups with tsinfer's own code. This step +needs a tsinfer checkout at `../../tsinfer`, and runs in its environment +so that tsinfer is not a dependency of the benchmarks: + +``` +uv run --project ../../tsinfer python tsinfer_groups.py data/tgp_chr20p +``` + +This writes `data/tgp_chr20p/match_jobs.csv`: 439,255 jobs in 1,487 +groups (groups 0–1485 for the ancestors, and group 1486 for the samples). +It checks that the order of the ancestor jobs is the node order of the +ancestors tree sequence, and that group 84 is the one in `group.json`. + +#### 2. Make workloads + +``` +uv run python tsinfer_data.py ancestors data/tgp_chr20p 1485 data/tgp_chr20p/ancestors_g1485 --num-queries 500 +uv run python tsinfer_data.py samples data/tgp_chr20p data/tgp_chr20p/samples --num-queries 200 +``` + +- `ancestors GROUP`: the ancestors of one group (all of them by default), + matched against the tree sequence before the group. This is the + ancestors tree sequence cut at the group's first node: tsinfer adds + each group's nodes in order, and a node only copies from earlier + groups. For group 84 the result is identical to `before.trees`. +- `samples`: a random subset of the sample haplotypes, matched against + the full ancestors tree sequence. The inference sites are selected from + the samples zarr with the config's include expression, then restricted + to the ancestor sites. Haplotypes are checked against the sample + genotypes in the matched tree sequence. + +Each workload directory holds `reference.trees`, the queries, a +`model.json` with tsinfer's parameters, and `tsinfer_group.json` with the +time and memory tsinfer recorded for the whole group. Haplotypes are +encoded as in tsinfer (0 ancestral, 1 derived, -1 missing), and each +query is matched over its window of non-missing sites, as tsinfer does. +Queries are stored in a windowed format (`queries.WindowedHaplotypes`), +because ancestors cover only a short stretch of the genome. +`queries.csv` has the path tsinfer found (`truth_parents`, +`truth_breakpoints`), taken from the edges of the query's node. + +#### 3. Run and summarise + +``` +uv run python run.py data/tgp_chr20p/ancestors_g1485/reference.trees \ + data/tgp_chr20p/ancestors_g1485 -o results/tgp_chr20p_ancestors_g1485.csv +uv run python tsinfer_summary.py \ + results/tgp_chr20p_ancestors_g1485.csv data/tgp_chr20p/ancestors_g1485 +``` + +When a workload has a `model.json`, `run.py` uses its parameters, runs +the `tsinfer` and `tsls` engines by default, and leaves `log_likelihood` +empty (the sc2ts formula does not apply with `weight_by_n=True`). +`tsinfer_summary.py` takes any number of `RESULTS QUERY_DIR` pairs, and +prints a markdown table. + +### Caveats + +- tskit-hmm cannot be used: the ancestors tree sequence has no samples, + and tskit-hmm can only copy from samples. +- tsinfer's recorded times are for whole groups, run with many threads + on a 196-core node. We report its CPU time per haplotype, which includes + reading genotypes and building the tree sequence, not just the HMM. +- `tsinfer_groups.py` uses tsinfer's private `config._parse_config` and + `pipeline._build_jobs_with_metadata`, because the config has no API for + changing its (absolute, cluster) paths. tsinfer's `match` with + `reference_ts` cannot be used to restart from a given tree sequence + ([tsinfer#1185](https://github.com/tskit-dev/tsinfer/issues/1185)), + which is why the workloads call the HMM directly. +- The local tsinfer checkout used for grouping was at `186e1af` + (a descendant of `cb5017e`); the groups are checked against the data + as described above. + + +### Findings so far (October 2026) + +Run one query at a time on a 4-core machine with 15 GB of RAM. Ancestor +groups 1483–1485 are random subsets of 500. tsinfer's time is its CPU time +for the whole group divided by the number of haplotypes. + +| workload | queries | mean sites | same path as tsinfer | tsinfer-hmm mean time | tsls mean time | tsinfer recorded | +|---|---|---|---|---|---|---| +| ancestors, group 84 | 144 of 144 | 1,390 | 100% (both) | 4.3 ms | 5.2 ms | 9.0 ms | +| ancestors, group 1000 | 70 of 70 | 2,950 | 100% (both) | 9.9 ms | 10.9 ms | 13.6 ms | +| ancestors, group 1483 | 500 of 25,935 | 12,200 | 100% (both) | 68 ms | 73 ms | 60 ms | +| ancestors, group 1484 | 500 of 40,684 | 62,000 | 100% (both) | 182 ms | 189 ms | 214 ms | +| ancestors, group 1485 | 500 of 79,127 | 40,200 | 100% (both) | 154 ms | 170 ms | 166 ms | +| samples | 200 of 6,404 | 580,909 | 100% (both) | 2.53 s | 2.80 s | 3.54 s | + +- Both engines find exactly the path tsinfer found for every query, so + they are drop-in replacements for tsinfer's matcher on this data. + Sample paths have no mismatches. Ancestor paths have one or two + mismatches on average: the ancestor's focal sites, whose mutations are + 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 + 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. diff --git a/benchmarking/engines.py b/benchmarking/engines.py index f105070..99b3274 100644 --- a/benchmarking/engines.py +++ b/benchmarking/engines.py @@ -103,41 +103,64 @@ def path_haplotype(ts, left, right, parent): return G[np.arange(ts.num_sites), index[segment]].astype(np.int8) -class TsinferEngine: +class AncestorMatcherEngine: """ + Base class for the engines with tsinfer's AncestorMatcher interface, + which differ only in how the MatcherIndexes are made. + If mismatch_threshold is specified, normalised likelihoods are not allowed to drop below that of mismatch_threshold mismatches, as in sc2ts. This bounds the number of distinct likelihood values, and so the time and memory needed, at the cost of exactness. - """ - name = "tsinfer" + 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. + """ - def __init__(self, ts, *, mu, rho, mismatch_threshold=None): + module = None + + def __init__( + self, + ts, + *, + mu, + rho, + mismatch_threshold=None, + weight_by_n=False, + alleles=ALLELES, + ): self.num_sites = ts.num_sites self.mu = mu self.rho = rho + self.weight_by_n = weight_by_n self.likelihood_threshold = None if mismatch_threshold is not None: 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 = tsinfer_hmm.MatcherIndexes( - ts, vestigial_root=False, alleles=ALLELES - ) + self.matcher_indexes = self.make_matcher_indexes(ts, alleles) - def match(self, h): - matcher = tsinfer_hmm.AncestorMatcher( + def make_matcher_indexes(self, ts, alleles): + raise NotImplementedError + + def match(self, h, start=0, end=None): + """ + Match the specified haplotype over the sites from start to end + (exclusive). The path covers only these sites. + """ + 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=False, + weight_by_n=self.weight_by_n, ) matched = np.zeros(self.num_sites, dtype=np.int8) - left, right, parent = matcher.find_path( - h.astype(np.int8), 0, self.num_sites, matched - ) + left, right, parent = matcher.find_path(h.astype(np.int8), start, end, matched) stats = { "total_memory": matcher.total_memory, "mean_traceback_size": matcher.mean_traceback_size, @@ -151,6 +174,14 @@ def match(self, h): ) +class TsinferEngine(AncestorMatcherEngine): + name = "tsinfer" + module = tsinfer_hmm + + def make_matcher_indexes(self, ts, alleles): + return tsinfer_hmm.MatcherIndexes(ts, vestigial_root=False, alleles=alleles) + + class TskitEngine: name = "tskit" @@ -203,55 +234,21 @@ def optimum_log_likelihood(self, normalisation_factor): return float(np.sum(np.log(normalisation_factor)) + np.log(n) - m * np.log(stay)) -class TslsEngine(TsinferEngine): +class TslsEngine(AncestorMatcherEngine): """ The merged tsls engine, with the same model and options as tsinfer. In tsls every node in the tree is a state, so we remove the vestigial root - edge 0 -> 1 added by prepare.py to make node 0 isolated, as in sc2ts. + edge 0 -> 1 (added by prepare.py, and present in tsinfer's trees) to + make node 0 isolated, as in sc2ts. """ name = "tsls" + module = tsls - def __init__(self, ts, *, mu, rho, mismatch_threshold=None): - self.num_sites = ts.num_sites - self.mu = mu - self.rho = rho - self.likelihood_threshold = None - if mismatch_threshold is not None: - self.likelihood_threshold = mu**mismatch_threshold * 0.99 - self.recombination = np.full(ts.num_sites, rho) - self.mismatch = np.full(ts.num_sites, mu) + def make_matcher_indexes(self, ts, alleles): tables = ts.dump_tables() tables.edges.keep_rows(tables.edges.parent != 0) - self.matcher_indexes = tsls.MatcherIndexes( - tables.tree_sequence(), alleles=ALLELES - ) - - def match(self, h): - matcher = tsls.AncestorMatcher( - self.matcher_indexes, - self.recombination, - self.mismatch, - likelihood_threshold=self.likelihood_threshold, - weight_by_n=False, - ) - matched = np.zeros(self.num_sites, dtype=np.int8) - left, right, parent = matcher.find_path( - h.astype(np.int8), 0, self.num_sites, matched - ) - stats = { - "total_memory": matcher.total_memory, - "mean_traceback_size": matcher.mean_traceback_size, - } - return ( - left[::-1].astype(np.float64), - right[::-1].astype(np.float64), - parent[::-1].copy(), - matched, - stats, - ) + return tsls.MatcherIndexes(tables.tree_sequence(), alleles=alleles) -ENGINES = { - engine.name: engine for engine in [TsinferEngine, TskitEngine, TslsEngine] -} +ENGINES = {engine.name: engine for engine in [TsinferEngine, TskitEngine, TslsEngine]} diff --git a/benchmarking/pyproject.toml b/benchmarking/pyproject.toml index 3a90092..80b7c19 100644 --- a/benchmarking/pyproject.toml +++ b/benchmarking/pyproject.toml @@ -13,6 +13,8 @@ dependencies = [ "tsinfer-hmm", "tskit-hmm", "tsls", + "zarr>=3.1", + "tabulate>=0.10.0", ] [dependency-groups] diff --git a/benchmarking/queries.py b/benchmarking/queries.py index d3ebe10..663d4a6 100644 --- a/benchmarking/queries.py +++ b/benchmarking/queries.py @@ -15,6 +15,7 @@ """ import dataclasses +import json import pathlib import click @@ -158,16 +159,105 @@ def save(self, out_dir): self.dataframe().to_csv(out_dir / "queries.csv", index=False) +class WindowedHaplotypes: + """ + Query haplotypes that are only defined over a window of sites, such as + tsinfer ancestors. Query j has the specified values at sites start[j] to + end[j] (exclusive) and is missing elsewhere. The values of all queries + are concatenated in a single array, so that short queries over many sites + do not need much space. + """ + + def __init__(self, num_sites, start, end, values): + self.num_sites = num_sites + self.start = np.asarray(start, dtype=np.int64) + self.end = np.asarray(end, dtype=np.int64) + self.values = np.asarray(values, dtype=np.int8) + lengths = self.end - self.start + assert np.all(lengths >= 0) + assert np.all(self.end <= num_sites) + self.offsets = np.concatenate([[0], np.cumsum(lengths)]) + assert self.offsets[-1] == len(self.values) + + @staticmethod + def from_windows(num_sites, start, windows): + """ + Return the WindowedHaplotypes made from the specified start sites and + list of window value arrays. + """ + start = np.asarray(start, dtype=np.int64) + end = start + np.array([len(w) for w in windows], dtype=np.int64) + values = np.concatenate([np.zeros(0, dtype=np.int8), *windows]) + return WindowedHaplotypes(num_sites, start, end, values) + + def __len__(self): + return len(self.start) + + def window(self, j): + return self.values[self.offsets[j] : self.offsets[j + 1]] + + def __getitem__(self, j): + h = np.full(self.num_sites, MISSING, dtype=np.int8) + h[self.start[j] : self.end[j]] = self.window(j) + return h + + def save(self, out_dir, positions, df): + assert len(df) == len(self) + assert len(positions) == self.num_sites + out_dir = pathlib.Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + out_dir / "queries.npz", + positions=positions, + start=self.start, + end=self.end, + values=self.values, + ) + df.to_csv(out_dir / "queries.csv", index=False) + + def load_queries(query_dir): """ Return the (H, positions, metadata) tuple for the query set in the - specified directory. + specified directory. H is either an array with one row per query, or + a WindowedHaplotypes; both give the haplotype of query j as H[j]. """ query_dir = pathlib.Path(query_dir) data = np.load(query_dir / "queries.npz") df = pd.read_csv(query_dir / "queries.csv") - assert len(df) == data["H"].shape[0] - return data["H"], data["positions"], df + positions = data["positions"] + if "values" in data: + H = WindowedHaplotypes( + len(positions), data["start"], data["end"], data["values"] + ) + else: + H = data["H"] + assert len(df) == len(H) + return H, positions, df + + +def query_window(H, query_id): + """ + Return the (start, end) tuple of the sites to match for the specified + query. + """ + if isinstance(H, WindowedHaplotypes): + return int(H.start[query_id]), int(H.end[query_id]) + return 0, H.shape[1] + + +def load_model(query_dir): + """ + Return the model parameters stored with the query set in the specified + directory, or None if there are none. Query sets made from tsinfer runs + store the parameters tsinfer used, which take the place of the sc2ts + style parameters chosen by run.py. + """ + path = pathlib.Path(query_dir) / "model.json" + if not path.exists(): + return None + with open(path) as f: + return json.load(f) def add_copy_queries(query_set, ts, rng, *, num_nodes, divergence): diff --git a/benchmarking/run.py b/benchmarking/run.py index 61b6fb2..78cb4f3 100644 --- a/benchmarking/run.py +++ b/benchmarking/run.py @@ -57,27 +57,31 @@ def completed_pairs(out_path): return set(zip(df.query_id, df.engine)) -def summarise_match(ts, node_ids, h, result, mu, rho): +def summarise_match(ts, node_ids, h, result, mu, rho, *, log_likelihood=True): """ Return the result row fields describing the specified match. The log likelihood is computed under the unweighted model with the specified - parameters, so that it is comparable between engines. Parents are - reported as original node IDs, so that they are comparable between - sub-ARGs. + parameters, so that it is comparable between engines; it is left empty + if log_likelihood is False. Parents are reported as original node IDs, + so that they are comparable between sub-ARGs. """ left, right, parent, matched, stats = result if matched is None: matched = engines.path_haplotype(ts, left, right, parent) non_missing = h != queries.MISSING - log_likelihood = engines.path_log_likelihood(h, len(parent), matched, mu, rho) - return { + row = { "num_segments": len(parent), "parents": queries.format_list(node_ids[parent]), "breakpoints": queries.format_list(left[1:]), "num_mismatches": int(np.sum(h[non_missing] != matched[non_missing])), - "log_likelihood": log_likelihood, + "log_likelihood": "", "stats": " ".join(f"{key}={value}" for key, value in stats.items()), } + if log_likelihood: + row["log_likelihood"] = engines.path_log_likelihood( + h, len(parent), matched, mu, rho + ) + return row @click.command() @@ -89,16 +93,14 @@ def summarise_match(ts, node_ids, h, result, mu, rho): "engine_names", type=click.Choice(list(engines.ENGINES)), multiple=True, - default=list(engines.ENGINES), - show_default=True, + help="Engines to run (default: all, or tsinfer and tsls with a model.json)", ) @click.option("-o", "--output", type=click.Path(dir_okay=False), required=True) @click.option( "--num-mismatches", type=int, - default=3, - show_default=True, - help="Number of mismatches equivalent to one recombination", + default=None, + help="Number of mismatches equivalent to one recombination [default: 3]", ) @click.option( "--mismatch-threshold", @@ -121,7 +123,32 @@ def main( """ Run the HMM engines on the queries in QUERY_DIR against the benchmark ARG at TS_PATH. + + If QUERY_DIR contains a model.json (as made by tsinfer_data.py), its + parameters are used instead of --num-mismatches, and the tskit engine, + which cannot use them, is not available. """ + model = queries.load_model(query_dir) + engine_options = {} + if model is None: + if len(engine_names) == 0: + engine_names = list(engines.ENGINES) + if num_mismatches is None: + num_mismatches = 3 + mu, rho = engines.solve_num_mismatches(num_mismatches) + else: + if len(engine_names) == 0: + engine_names = ["tsinfer", "tsls"] + if "tskit" in engine_names: + raise click.UsageError("The tskit engine does not support model.json") + if num_mismatches is not None: + raise click.UsageError("--num-mismatches cannot be used with model.json") + mu = model["mu"] + rho = model["rho"] + engine_options = { + "weight_by_n": model["weight_by_n"], + "alleles": model["alleles"], + } ts = load_ts(ts_path) node_ids = subargs.load_node_ids(ts_path, ts) H, positions, df = queries.load_queries(query_dir) @@ -137,7 +164,6 @@ def main( done = completed_pairs(out_path) write_header = len(done) == 0 - mu, rho = engines.solve_num_mismatches(num_mismatches) with open(out_path, "a", newline="") as f: writer = csv.DictWriter(f, fieldnames=RESULT_COLUMNS) if write_header: @@ -150,11 +176,13 @@ def main( options = {} if name in ("tsinfer", "tsls"): options["mismatch_threshold"] = mismatch_threshold + options.update(engine_options) engine = engines.ENGINES[name](ts, mu=mu, rho=rho, **options) setup_time = time.perf_counter() - before click.echo(f"Set up {name} in {setup_time:.2f}s") for query_id in tqdm(todo, desc=name): h = H[query_id] + start, end = queries.query_window(H, query_id) row = { "query_id": query_id, "engine": name, @@ -167,7 +195,10 @@ def main( wall_before = time.perf_counter() cpu_before = time.thread_time() try: - result = engine.match(h) + if name == "tskit": + result = engine.match(h) + else: + result = engine.match(h, start, end) except Exception as e: # noqa: BLE001 (record any failure) row["status"] = f"{type(e).__name__}: {e}" else: @@ -175,7 +206,17 @@ def main( row["wall_time"] = time.perf_counter() - wall_before row["cpu_time"] = time.thread_time() - cpu_before if row["status"] == "ok": - row.update(summarise_match(ts, node_ids, h, result, mu, rho)) + row.update( + summarise_match( + ts, + node_ids, + h, + result, + mu, + rho, + log_likelihood=model is None, + ) + ) writer.writerow(row) f.flush() diff --git a/benchmarking/tests/test_tsinfer_data.py b/benchmarking/tests/test_tsinfer_data.py new file mode 100644 index 0000000..c8c2592 --- /dev/null +++ b/benchmarking/tests/test_tsinfer_data.py @@ -0,0 +1,322 @@ +import json +import pathlib + +import msprime +import numpy as np +import pandas as pd +import pytest +import tskit +from click.testing import CliRunner + +import engines +import queries +import run +import tsinfer_data + +DATA_DIR = pathlib.Path(__file__).parent.parent / "data" / "tgp_chr20p" +TSINFER_MODEL = {"rho": 0.01, "mu": 1e-20, "weight_by_n": True, "alleles": None} + + +def simulate_tsinfer_like(seed, samples=20): + """ + Simulate a haploid ARG with the structure of a tsinfer ancestors tree + sequence: nodes in decreasing time order, starting with the ultimate + root 0 and the virtual root 1, which is the root of every tree, and + biallelic sites with a single 0 -> 1 mutation each. + """ + ts = msprime.sim_ancestry( + samples, + ploidy=1, + sequence_length=1000, + recombination_rate=1e-3, + random_seed=seed, + ) + ts = msprime.sim_mutations( + ts, + rate=5e-3, + model=msprime.BinaryMutationModel(), + random_seed=seed, + ) + tables = ts.dump_tables() + num_mutations = np.bincount(tables.mutations.site, minlength=ts.num_sites) + tables.delete_sites(np.where(num_mutations != 1)[0]) + max_time = np.max(tables.nodes.time) + tables.nodes.clear() + tables.nodes.add_row(time=max_time + 2) + tables.nodes.add_row(time=max_time + 1) + for node in ts.tables.nodes: + tables.nodes.append(node.replace(flags=0)) + tables.edges.parent += 2 + tables.edges.child += 2 + tables.mutations.node += 2 + for tree in ts.trees(): + for root in tree.roots: + tables.edges.add_row(*tree.interval, parent=1, child=root + 2) + tables.edges.add_row(0, ts.sequence_length, parent=0, child=1) + tables.edges.squash() + tables.sort() + ts = tables.tree_sequence() + order = np.argsort(-ts.nodes_time, kind="stable") + ts = ts.subset(order) + assert ts.num_sites > 20 + assert np.all(np.diff(ts.nodes_time) <= 0) + return ts + + +@pytest.fixture(scope="module") +def ts(): + return simulate_tsinfer_like(42) + + +def node_haplotypes(ts, nodes): + return ts.genotype_matrix(samples=nodes).T.astype(np.int8) + + +class TestCutBefore: + @pytest.mark.parametrize("cut", [2, 10, 25]) + def test_haplotypes_unchanged(self, ts, cut): + sub = tsinfer_data.cut_before(ts, cut) + assert sub.num_nodes == cut + assert sub.num_sites == ts.num_sites + nodes = np.arange(cut) + np.testing.assert_array_equal( + node_haplotypes(sub, nodes), node_haplotypes(ts, nodes) + ) + + def test_full_cut_is_identity(self, ts): + sub = tsinfer_data.cut_before(ts, ts.num_nodes) + sub.tables.assert_equals(ts.tables) + + def test_later_parent_fails(self, ts): + # Youngest nodes first, so children come before their parents. + reversed_ts = ts.subset(np.arange(ts.num_nodes)[::-1]) + with pytest.raises(ValueError, match="later node"): + tsinfer_data.cut_before(reversed_ts, ts.num_nodes // 2) + + +class TestTruthPaths: + def test_paths_cover_edges(self, ts): + nodes = np.arange(2, ts.num_nodes) + paths = tsinfer_data.truth_paths(ts, nodes) + for node, (left, parent) in zip(nodes, paths): + edges = ts.tables.edges + keep = edges.child == node + assert len(left) == np.sum(keep) + assert np.all(np.diff(left) > 0) + assert set(zip(left, parent)) == set( + zip(edges.left[keep], edges.parent[keep]) + ) + + def test_columns(self): + paths = [(np.array([0.0, 5.0]), np.array([3, 4])), (np.array([2.0]), [1])] + columns = tsinfer_data.truth_columns(paths) + assert columns["truth_num_segments"] == [2, 1] + assert columns["truth_parents"] == ["3 4", "1"] + assert columns["truth_breakpoints"] == ["5", ""] + + +class TestWindowedHaplotypes: + def make(self): + windows = [np.array([0, 1, 1]), np.array([], dtype=np.int8), np.array([1])] + return queries.WindowedHaplotypes.from_windows(6, [1, 0, 5], windows) + + def test_getitem(self): + H = self.make() + assert len(H) == 3 + np.testing.assert_array_equal(H[0], [-1, 0, 1, 1, -1, -1]) + np.testing.assert_array_equal(H[1], [-1] * 6) + np.testing.assert_array_equal(H[2], [-1, -1, -1, -1, -1, 1]) + assert queries.query_window(H, 0) == (1, 4) + assert queries.query_window(H, 2) == (5, 6) + + def test_dense_window(self): + H = np.zeros((2, 7), dtype=np.int8) + assert queries.query_window(H, 1) == (0, 7) + + def test_save_load(self, tmp_path): + H = self.make() + df = pd.DataFrame({"query_id": np.arange(3)}) + H.save(tmp_path, np.arange(6) * 10, df) + H2, positions, df2 = queries.load_queries(tmp_path) + np.testing.assert_array_equal(positions, np.arange(6) * 10) + for j in range(3): + np.testing.assert_array_equal(H[j], H2[j]) + assert len(df2) == 3 + + def test_non_missing_window(self): + assert tsinfer_data.non_missing_window(np.array([-1, 0, -1, 1, -1])) == (1, 4) + assert tsinfer_data.non_missing_window(np.array([-1, -1])) == (0, 2) + + +@pytest.mark.parametrize("name", ["tsinfer", "tsls"]) +class TestTsinferModel: + def make_engine(self, ts, name): + return engines.ENGINES[name]( + ts, + mu=TSINFER_MODEL["mu"], + rho=TSINFER_MODEL["rho"], + weight_by_n=True, + alleles=None, + ) + + @pytest.mark.parametrize("window", [(0, None), (5, 15), (10, 11)]) + def test_exact_copy(self, ts, name, window): + engine = self.make_engine(ts, name) + start, end = window + if end is None: + end = ts.num_sites + positions = ts.sites_position + for node in range(ts.num_nodes - 5, ts.num_nodes): + h = np.full(ts.num_sites, -1, dtype=np.int8) + h[start:end] = node_haplotypes(ts, [node])[0, start:end] + left, right, _, matched, _ = engine.match(h, start, end) + # As in tsinfer, a path starting at the first site starts at 0. + assert left[0] == (0 if start == 0 else positions[start]) + if end < ts.num_sites: + assert right[-1] == positions[end] + else: + assert right[-1] == ts.sequence_length + np.testing.assert_array_equal(left[1:], right[:-1]) + assert np.all(matched[start:end] == h[start:end]) + + def test_engines_agree_on_mismatches(self, ts, name): + engine = self.make_engine(ts, name) + other = self.make_engine(ts, "tsinfer") + rng = np.random.default_rng(1) + for node in range(ts.num_nodes - 5, ts.num_nodes): + h = node_haplotypes(ts, [node])[0] + flip = rng.choice(ts.num_sites, size=3, replace=False) + h[flip] = 1 - h[flip] + matched = engine.match(h)[3] + other_matched = other.match(h)[3] + assert np.sum(matched != h) == np.sum(other_matched != h) + + +class TestRunWithModel: + @pytest.fixture + def workload(self, ts, tmp_path): + reference = tsinfer_data.cut_before(ts, ts.num_nodes - 5) + nodes = np.arange(ts.num_nodes - 5, ts.num_nodes) + G = node_haplotypes(ts, nodes) + starts = [0, 3, 10, 0, 7] + windows = [G[j, s : s + 10] for j, s in enumerate(starts)] + H = queries.WindowedHaplotypes.from_windows(ts.num_sites, starts, windows) + df = pd.DataFrame( + { + "query_id": np.arange(len(nodes)), + "kind": "ancestor", + "sample_id": [f"a{u}" for u in nodes], + "ploidy_index": 0, + "node": nodes, + "group": 0, + "start": H.start, + "end": H.end, + **tsinfer_data.truth_columns(tsinfer_data.truth_paths(ts, nodes)), + } + ) + resources = {"group": 0, "num_haplotypes": 5, "user_time": 1.0} + tsinfer_data.save_workload( + tmp_path / "w", reference, H, df, TSINFER_MODEL, resources + ) + return tmp_path / "w" + + def test_run(self, workload, tmp_path): + out = tmp_path / "results.csv" + result = CliRunner().invoke( + run.main, [str(workload / "reference.trees"), str(workload), "-o", str(out)] + ) + assert result.exit_code == 0, result.output + df = pd.read_csv(out) + assert list(df.engine.unique()) == ["tsinfer", "tsls"] + assert np.all(df.status == "ok") + assert np.all(df.rho == 0.01) + assert np.all(df.log_likelihood.isna()) + H, positions, _ = queries.load_queries(workload) + for row in df.itertuples(): + start, end = queries.query_window(H, row.query_id) + breakpoints = queries.parse_list(row.breakpoints) + assert len(breakpoints) == row.num_segments - 1 + assert np.all(np.isin(breakpoints, positions[start + 1 : end])) + + def test_tskit_rejected(self, workload, tmp_path): + result = CliRunner().invoke( + run.main, + [ + str(workload / "reference.trees"), + str(workload), + "-o", + str(tmp_path / "out.csv"), + "-e", + "tskit", + ], + ) + assert result.exit_code != 0 + assert "model.json" in result.output + + def test_num_mismatches_rejected(self, workload, tmp_path): + result = CliRunner().invoke( + run.main, + [ + str(workload / "reference.trees"), + str(workload), + "-o", + str(tmp_path / "out.csv"), + "--num-mismatches", + "4", + ], + ) + assert result.exit_code != 0 + + def test_load_model(self, workload, tmp_path): + assert queries.load_model(workload) == TSINFER_MODEL + assert queries.load_model(tmp_path) is None + with open(workload / "tsinfer_group.json") as f: + assert json.load(f)["group"] == 0 + + +needs_data = pytest.mark.skipif( + not (DATA_DIR / tsinfer_data.MATCH_JOBS).exists(), + reason="Needs data/tgp_chr20p with match_jobs.csv (see README)", +) + + +@pytest.fixture(scope="module") +def workload_g84(): + return tsinfer_data.make_ancestor_workload(DATA_DIR, 84, num_queries=10) + + +@needs_data +class TestTgpChr20p: + """ + Checks against the partial ancestor match supplied with the data: the + tree sequence before group 84 and the paths that tsinfer found for it. + """ + + @pytest.fixture + def workload(self, workload_g84): + return workload_g84 + + def test_reference_is_before(self, workload): + reference = workload[0] + before = tskit.load(DATA_DIR / "partial_ancestor_match" / "before.trees") + reference.tables.assert_equals(before.tables, ignore_provenance=True) + + def test_resources(self, workload): + resources = workload[3] + assert resources["group"] == 84 + assert resources["num_haplotypes"] == 144 + + @pytest.mark.parametrize("name", ["tsinfer", "tsls"]) + def test_paths_match_tsinfer(self, workload, name): + reference, H, df, _ = workload + after = tskit.load(DATA_DIR / "partial_ancestor_match" / "after.trees") + engine = engines.ENGINES[name]( + reference, mu=1e-20, rho=0.01, weight_by_n=True, alleles=None + ) + for row in df.itertuples(): + assert after.node(row.node).metadata["sample_id"] == row.sample_id + start, end = queries.query_window(H, row.query_id) + left, _, parent, _, _ = engine.match(H[row.query_id], start, end) + ((truth_left, truth_parent),) = tsinfer_data.truth_paths(after, [row.node]) + np.testing.assert_array_equal(left, truth_left) + np.testing.assert_array_equal(parent, truth_parent) diff --git a/benchmarking/tsinfer_data.py b/benchmarking/tsinfer_data.py new file mode 100644 index 0000000..dc0100c --- /dev/null +++ b/benchmarking/tsinfer_data.py @@ -0,0 +1,372 @@ +""" +Make benchmark workloads from a tsinfer 1.0 inference. + +The data directory holds the inputs and outputs of a tsinfer run (see +data/tgp_chr20p/README.md), plus the match_jobs.csv file written by +tsinfer_groups.py. There are two kinds of workload, matching the two +phases of tsinfer's matching: + +- ancestors: the ancestors of one match group, matched against the tree + sequence made from all earlier groups; +- samples: sample haplotypes, matched against the full ancestors tree + sequence. + +Each workload is a directory that run.py can use, containing: + +- reference.trees: the tree sequence to match against; +- queries.npz, queries.csv: the query haplotypes, in the windowed format + (see queries.WindowedHaplotypes), and their metadata, including the path + that tsinfer found (truth_parents and truth_breakpoints); +- model.json: the HMM parameters used by tsinfer; +- tsinfer_group.json: the resources tsinfer recorded for matching the + whole group (elapsed and CPU times are for all threads). + +Haplotypes are encoded as in tsinfer: 0 for the ancestral allele, 1 for +the derived allele and -1 for missing data. +""" + +import json +import pathlib +import tomllib + +import click +import numpy as np +import pandas as pd +import tskit +import zarr + +import queries + +ANCESTORS_TS = "tgp_chr20p_inferred_ancestors.trees" +ANCESTORS_ZARR = "tgp_chr20p_inferred_ancestors.zarr" +SAMPLES_TS = "tgp_chr20p_inferred_samples_raw.trees" +SAMPLES_ZARR = "tgp_chr20p_samples_masked.zarr" +CONFIG = "full_inference.toml" +MATCH_JOBS = "match_jobs.csv" + +# Nodes 0 and 1 of a tsinfer tree sequence are the ultimate and virtual +# roots, so the node of the jth matched haplotype is j + 2. +NUM_ROOTS = 2 + +QUERY_COLUMNS = [ + "query_id", + "kind", + "sample_id", + "ploidy_index", + "node", + "group", + "start", + "end", + "truth_num_segments", + "truth_parents", + "truth_breakpoints", +] + + +def cut_before(ts, node): + """ + Return the tree sequence made of the nodes before the specified node ID, + with their edges and mutations. Because tsinfer adds the nodes of each + group in order, and a node can only copy from earlier groups, cutting + before the first node of a group gives the tree sequence that the group + was matched against. + """ + tables = ts.dump_tables() + tables.edges.keep_rows(tables.edges.child < node) + if np.any(tables.edges.parent >= node): + raise ValueError("Edge from a later node to an earlier one") + tables.mutations.keep_rows(tables.mutations.node < node) + tables.nodes.truncate(node) + tables.build_index() + return tables.tree_sequence() + + +def truth_paths(ts, nodes): + """ + Return the list of (left, parent) tuples giving the path of each of the + specified nodes, as arrays in left-to-right order. + """ + edges = ts.tables.edges + order = np.lexsort((edges.left, edges.child)) + child = edges.child[order] + lo = np.searchsorted(child, nodes, side="left") + hi = np.searchsorted(child, nodes, side="right") + return [(edges.left[order[a:b]], edges.parent[order[a:b]]) for a, b in zip(lo, hi)] + + +def truth_columns(paths): + return { + "truth_num_segments": [len(parent) for _, parent in paths], + "truth_parents": [queries.format_list(parent) for _, parent in paths], + "truth_breakpoints": [queries.format_list(left[1:]) for left, _ in paths], + } + + +def non_missing_window(h): + """ + Return the (start, end) tuple of the first and one past the last + non-missing site, as tsinfer matches over this range. + """ + non_missing = np.where(h != queries.MISSING)[0] + if len(non_missing) == 0: + return 0, len(h) + return int(non_missing[0]), int(non_missing[-1]) + 1 + + +def load_model(data_dir, source): + """ + Return the HMM parameters that tsinfer used for the specified source, + with tsinfer's defaults where the config does not set them. + """ + with open(data_dir / CONFIG, "rb") as f: + raw = tomllib.load(f) + params = raw["match"]["sources"][source] + return { + "rho": params.get("recombination", 1e-2), + "mu": params.get("mismatch", 1e-20), + "weight_by_n": True, + "alleles": None, + } + + +def group_resources(ts, group, group_size): + """ + Return the resources that tsinfer recorded in the provenance for + matching the specified group, with the number of haplotypes in it. + """ + for provenance in ts.provenances(): + record = json.loads(provenance.record) + parameters = record.get("parameters", {}) + if parameters.get("command") == "match" and parameters.get("group") == group: + return {"group": group, "num_haplotypes": group_size, **record["resources"]} + raise ValueError(f"No provenance for group {group}") + + +def save_workload(out_dir, reference, H, df, model, resources): + out_dir = pathlib.Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + reference.dump(out_dir / "reference.trees") + H.save(out_dir, reference.sites_position, df[QUERY_COLUMNS]) + with open(out_dir / "model.json", "w") as f: + json.dump(model, f, indent=2) + with open(out_dir / "tsinfer_group.json", "w") as f: + json.dump(resources, f, indent=2) + + +def choose(rng, n, num_queries): + """ + Return the sorted indexes of num_queries of n items chosen at random, or + all of them if num_queries is None. + """ + if num_queries is None or num_queries >= n: + return np.arange(n) + return np.sort(rng.choice(n, size=num_queries, replace=False)) + + +def load_jobs(data_dir): + jobs = pd.read_csv(data_dir / MATCH_JOBS) + # The nodes are added in the order of the jobs. + jobs["node"] = np.arange(len(jobs)) + NUM_ROOTS + return jobs + + +def make_ancestor_workload(data_dir, group, *, num_queries=None, seed=42): + """ + Return the (reference, H, df, resources) tuple for the ancestors of the + specified group. + """ + jobs = load_jobs(data_dir) + jobs = jobs[jobs.source == "ancestors"] + group_jobs = jobs[jobs.group == group] + if len(group_jobs) == 0: + raise ValueError(f"No ancestors in group {group}") + full = tskit.load(data_dir / ANCESTORS_TS) + reference = cut_before(full, group_jobs.node.min()) + resources = group_resources(full, group, len(group_jobs)) + + rng = np.random.default_rng(seed) + group_jobs = group_jobs.iloc[choose(rng, len(group_jobs), num_queries)] + store = zarr.open_group(data_dir / ANCESTORS_ZARR, mode="r") + positions = store["variant_position"][:] + assert np.array_equal(positions, full.sites_position) + call_genotype = store["call_genotype"] + + starts = [] + windows = [] + for job in group_jobs.itertuples(): + # The ancestor zarr has one column per ancestor, in job order. + assert job.sample_id == f"a{job.haplotype_index}" + lo = np.searchsorted(positions, job.start_position, side="left") + hi = np.searchsorted(positions, job.end_position, side="left") + h = call_genotype[lo:hi, job.haplotype_index, 0].astype(np.int8) + start, end = non_missing_window(h) + starts.append(lo + start) + windows.append(h[start:end]) + H = queries.WindowedHaplotypes.from_windows(len(positions), starts, windows) + + nodes = group_jobs.node.to_numpy() + df = pd.DataFrame( + { + "query_id": np.arange(len(group_jobs)), + "kind": "ancestor", + "sample_id": group_jobs.sample_id.to_numpy(), + "ploidy_index": group_jobs.ploidy_index.to_numpy(), + "node": nodes, + "group": group, + "start": H.start, + "end": H.end, + **truth_columns(truth_paths(full, nodes)), + } + ) + return reference, H, df, resources + + +def inference_sites(samples_store, ancestor_positions): + """ + Return the indexes of the inference sites in the samples store, and + their ancestral allele indexes. The sites are those selected by the + include expression in the config, + + (POS < 28100000) & ((INFO/match_eval_inference_site_mask == 0) + & (INFO/match_eval_derived_ac != 1)) + + that were used for ancestor generation, which drops some more sites. + The samples store has duplicate positions, but none of them survive + the selection. + """ + position = samples_store["variant_position"][:] + include = ( + (position < 28100000) + & (samples_store["variant_match_eval_inference_site_mask"][:] == 0) + & (samples_store["variant_match_eval_derived_ac"][:] != 1) + ) + selected = np.where(include)[0] + selected = selected[np.isin(position[selected], ancestor_positions)] + if not np.array_equal(position[selected], ancestor_positions): + raise ValueError("Inference sites do not match the ancestor sites") + alleles = samples_store["variant_allele"].get_orthogonal_selection((selected,)) + ancestral = samples_store["variant_ancestral_allele"].get_orthogonal_selection( + (selected,) + ) + ancestral_index = np.full(len(selected), -1, dtype=np.int8) + for j in range(alleles.shape[1]): + ancestral_index[alleles[:, j] == ancestral] = j + if np.any(ancestral_index < 0): + raise ValueError("Ancestral allele not found") + return selected, ancestral_index + + +def make_sample_workload(data_dir, *, num_queries=200, seed=42): + """ + Return the (reference, H, df, resources) tuple for a random subset of + the sample haplotypes. + """ + jobs = load_jobs(data_dir) + jobs = jobs[jobs.source == "samples"] + raw = tskit.load(data_dir / SAMPLES_TS) + (group,) = np.unique(jobs.group) + resources = group_resources(raw, int(group), len(jobs)) + assert np.array_equal(raw.samples(), jobs.node.to_numpy()) + rng = np.random.default_rng(seed) + jobs = jobs.iloc[choose(rng, len(jobs), num_queries)] + for job in jobs.itertuples(): + metadata = raw.node(job.node).metadata + assert metadata["sample_id"] == job.sample_id + assert metadata["ploidy_index"] == job.ploidy_index + + reference = tskit.load(data_dir / ANCESTORS_TS) + store = zarr.open_group(data_dir / SAMPLES_ZARR, mode="r") + sites, ancestral_index = inference_sites(store, reference.sites_position) + sample_ids = list(store["sample_id"][:]) + individuals = np.array([sample_ids.index(sid) for sid in jobs.sample_id]) + unique_individuals, inverse = np.unique(individuals, return_inverse=True) + G = store["call_genotype"].get_orthogonal_selection( + (sites, unique_individuals, slice(None)) + ) + mask = store["call_genotype_mask"].get_orthogonal_selection( + (sites, unique_individuals, slice(None)) + ) + G = G[:, inverse, jobs.ploidy_index.to_numpy()] + mask = mask[:, inverse, jobs.ploidy_index.to_numpy()] + H = (G != ancestral_index[:, np.newaxis]).astype(np.int8) + H[(G < 0) | mask] = queries.MISSING + H = H.T + + # The matched tree sequence has no mutations above the samples, so the + # sample genotypes are the input haplotypes at non-missing sites. + nodes = jobs.node.to_numpy() + expected = raw.genotype_matrix(samples=nodes).T + non_missing = H != queries.MISSING + if not np.array_equal(H[non_missing], expected[non_missing]): + raise ValueError("Sample haplotypes do not match the matched tree sequence") + + starts = [] + windows = [] + for h in H: + start, end = non_missing_window(h) + starts.append(start) + windows.append(h[start:end]) + H = queries.WindowedHaplotypes.from_windows(reference.num_sites, starts, windows) + df = pd.DataFrame( + { + "query_id": np.arange(len(jobs)), + "kind": "sample", + "sample_id": jobs.sample_id.to_numpy(), + "ploidy_index": jobs.ploidy_index.to_numpy(), + "node": nodes, + "group": jobs.group.to_numpy(), + "start": H.start, + "end": H.end, + **truth_columns(truth_paths(raw, nodes)), + } + ) + return reference, H, df, resources + + +@click.group() +def cli(): + pass + + +@cli.command() +@click.argument("data_dir", type=click.Path(exists=True, file_okay=False)) +@click.argument("group", type=int) +@click.argument("out_dir", type=click.Path(file_okay=False)) +@click.option("--num-queries", type=int, default=None, help="[default: all]") +@click.option("--seed", type=int, default=42, show_default=True) +def ancestors(data_dir, group, out_dir, num_queries, seed): + """ + Make the workload for the ancestors in GROUP. + """ + data_dir = pathlib.Path(data_dir) + reference, H, df, resources = make_ancestor_workload( + data_dir, group, num_queries=num_queries, seed=seed + ) + model = load_model(data_dir, "ancestors") + save_workload(out_dir, reference, H, df, model, resources) + click.echo( + f"Wrote {len(df)} ancestors (mean {np.mean(df.end - df.start):.0f} sites) " + f"and a reference with {reference.num_nodes} nodes to {out_dir}" + ) + + +@cli.command() +@click.argument("data_dir", type=click.Path(exists=True, file_okay=False)) +@click.argument("out_dir", type=click.Path(file_okay=False)) +@click.option("--num-queries", type=int, default=200, show_default=True) +@click.option("--seed", type=int, default=42, show_default=True) +def samples(data_dir, out_dir, num_queries, seed): + """ + Make the workload for a random subset of the sample haplotypes. + """ + data_dir = pathlib.Path(data_dir) + reference, H, df, resources = make_sample_workload( + data_dir, num_queries=num_queries, seed=seed + ) + model = load_model(data_dir, "samples") + save_workload(out_dir, reference, H, df, model, resources) + click.echo(f"Wrote {len(df)} sample haplotypes to {out_dir}") + + +if __name__ == "__main__": + cli() diff --git a/benchmarking/tsinfer_groups.py b/benchmarking/tsinfer_groups.py new file mode 100644 index 0000000..048d02b --- /dev/null +++ b/benchmarking/tsinfer_groups.py @@ -0,0 +1,101 @@ +""" +Compute tsinfer's match groups for the 1000 Genomes chr20p data. + +The ancestors tree sequence does not record which nodes were matched in +which group, so we recompute the groups with tsinfer's own grouping code. +This script must be run in the tsinfer environment, which keeps tsinfer +(GPL) out of the benchmarking dependencies: + + uv run --project ../../tsinfer python tsinfer_groups.py data/tgp_chr20p + +It writes match_jobs.csv to the data directory, with one row per matched +haplotype (ancestors then samples), ordered by group then haplotype index +as in tsinfer. The order is checked against the node order of the +ancestors tree sequence, and group 84 against the group.json file from +the partial ancestor match. +""" + +import json +import pathlib +import sys +import tomllib + +import numpy as np +import pandas as pd +import tskit +from tsinfer import config, pipeline + +COLUMNS = [ + "haplotype_index", + "source", + "sample_id", + "ploidy_index", + "time", + "start_position", + "end_position", + "group", +] + + +def load_config(data_dir): + """ + Return the tsinfer Config for the inference, with paths pointing to the + local copies of the stores. The toml file has absolute paths on the + cluster where it was run, and tsinfer's config has no API for changing + them, so we change the raw dict and use the private parser. + """ + with open(data_dir / "full_inference.toml", "rb") as f: + raw = tomllib.load(f) + samples_path = str(data_dir / "tgp_chr20p_samples_masked.zarr") + ancestors_path = str(data_dir / "tgp_chr20p_inferred_ancestors.zarr") + for source in raw["source"]: + assert source["name"] == "samples" + source["path"] = samples_path + raw["ancestral_state"]["path"] = samples_path + for ancestors in raw["ancestors"]: + ancestors["path"] = ancestors_path + raw["match"]["output"] = str(data_dir / "unused.trees") + return config._parse_config(raw) + + +def check_ancestor_order(jobs, ts): + """ + The nodes of the ancestors tree sequence are the two roots followed by + the ancestors in match order, so their sample IDs must be in the same + order as the ancestor jobs. + """ + ancestors = jobs[jobs.source == "ancestors"] + assert ts.num_nodes == len(ancestors) + 2 + node_ids = [ts.node(u).metadata["sample_id"] for u in range(2, ts.num_nodes)] + if node_ids != list(ancestors.sample_id): + raise ValueError("Ancestor job order does not match tree sequence node order") + + +def check_group(jobs, group_json_path): + with open(group_json_path) as f: + group = json.load(f) + expected = [job["haplotype_index"] for job in group["jobs"]] + actual = list(jobs.haplotype_index[jobs.group == group["group"]]) + if actual != expected: + raise ValueError(f"Group {group['group']} does not match {group_json_path}") + + +def main(data_dir): + data_dir = pathlib.Path(data_dir).resolve() + cfg = load_config(data_dir) + jobs = pipeline._build_jobs_with_metadata(cfg)[0] + df = pd.DataFrame( + {column: [getattr(job, column) for job in jobs] for column in COLUMNS} + ) + ts = tskit.load(data_dir / "tgp_chr20p_inferred_ancestors.trees") + check_ancestor_order(df, ts) + check_group(df, data_dir / "partial_ancestor_match" / "group.json") + sample_groups = np.unique(df.group[df.source == "samples"]) + assert len(sample_groups) == 1 + out_path = data_dir / "match_jobs.csv" + df.to_csv(out_path, index=False) + print(f"Wrote {len(df)} jobs in {df.group.nunique()} groups to {out_path}") + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/benchmarking/tsinfer_summary.py b/benchmarking/tsinfer_summary.py new file mode 100644 index 0000000..3ad8f2d --- /dev/null +++ b/benchmarking/tsinfer_summary.py @@ -0,0 +1,96 @@ +""" +Summarise run.py results on workloads made by tsinfer_data.py. + +For each workload and engine, report how often the engine found the same +path as tsinfer, along with its time and memory, and the time tsinfer +recorded for matching the whole group. +""" + +import json +import pathlib + +import click +import numpy as np +import pandas as pd + + +def parse_stats(stats): + """ + Return the dict of engine statistics in a stats column value. + """ + items = (item.split("=") for item in str(stats).split()) + return {key: float(value) for key, value in items} + + +def join_truth(results, query_dir): + """ + Return the results joined with the query metadata, with a same_path + column saying whether the engine's path is the one tsinfer found. + """ + df = pd.read_csv(query_dir / "queries.csv") + df = results.merge(df, on="query_id", how="left", validate="many_to_one") + + def column(name): + return df[name].fillna("").astype(str) + + df["same_path"] = (column("parents") == column("truth_parents")) & ( + column("breakpoints") == column("truth_breakpoints") + ) + stats = pd.DataFrame([parse_stats(s) for s in df.stats], index=df.index) + return pd.concat([df, stats], axis=1) + + +def summarise(results_path, query_dir): + query_dir = pathlib.Path(query_dir) + df = join_truth(pd.read_csv(results_path), query_dir) + with open(query_dir / "tsinfer_group.json") as f: + group = json.load(f) + rows = [] + for engine, d in df.groupby("engine", sort=False): + ok = d[d.status == "ok"] + rows.append( + { + "workload": query_dir.name, + "engine": engine, + "queries": len(d), + "failed": int(np.sum(d.status != "ok")), + "mean sites": np.mean(ok.end - ok.start), + "same path": np.mean(ok.same_path), + "segment diff": np.mean(ok.num_segments - ok.truth_num_segments), + "mean mismatches": np.mean(ok.num_mismatches), + "mean time (s)": np.mean(ok.wall_time), + "max time (s)": np.max(ok.wall_time), + "mean memory (MiB)": np.mean(ok.total_memory) / 2**20, + "mean traceback": np.mean(ok.mean_traceback_size), + "group size": group["num_haplotypes"], + "tsinfer CPU time per haplotype (s)": group["user_time"] + / group["num_haplotypes"], + } + ) + return pd.DataFrame(rows) + + +@click.command() +@click.argument( + "pairs", + nargs=-1, + required=True, + type=click.Path(exists=True), + metavar="RESULTS QUERY_DIR ...", +) +def main(pairs): + """ + Summarise each RESULTS file from run.py on the workload in QUERY_DIR, + as a markdown table. + """ + if len(pairs) % 2 != 0: + raise click.UsageError("Specify RESULTS QUERY_DIR pairs") + tables = [ + summarise(results_path, query_dir) + for results_path, query_dir in zip(pairs[::2], pairs[1::2]) + ] + click.echo(pd.concat(tables).to_markdown(index=False, floatfmt=".3g")) + + +if __name__ == "__main__": + main() diff --git a/benchmarking/uv.lock b/benchmarking/uv.lock index e7f104a..86bb6a1 100644 --- a/benchmarking/uv.lock +++ b/benchmarking/uv.lock @@ -852,6 +852,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] +[[package]] +name = "tabulate" +version = "0.10.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/46/58/8c37dea7bbf769b20d58e7ace7e5edfe65b849442b00ffcdd56be88697c6/tabulate-0.10.0.tar.gz", hash = "sha256:e2cfde8f79420f6deeffdeda9aaec3b6bc5abce947655d17ac662b126e48a60d", size = 91754, upload-time = "2026-03-04T18:55:34.402Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/55/db07de81b5c630da5cbf5c7df646580ca26dfaefa593667fc6f2fe016d2e/tabulate-0.10.0-py3-none-any.whl", hash = "sha256:f0b0622e567335c8fabaaa659f1b33bcb6ddfe2e496071b743aa113f8774f2d3", size = 39814, upload-time = "2026-03-04T18:55:31.284Z" }, +] + [[package]] name = "tqdm" version = "4.70.1" @@ -965,7 +974,10 @@ requires-dist = [ ] [package.metadata.requires-dev] -test = [{ name = "pytest" }] +test = [ + { name = "pytest" }, + { name = "pytest-xdist", specifier = ">=3.8.0" }, +] [[package]] name = "tsls-benchmarking" @@ -976,12 +988,15 @@ dependencies = [ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "numpy", version = "2.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "pandas" }, + { name = "tabulate" }, { name = "tqdm" }, { name = "tsinfer-hmm" }, { name = "tskit" }, { name = "tskit-hmm" }, { name = "tsls" }, { name = "tszip" }, + { name = "zarr", version = "3.1.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "zarr", version = "3.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] [package.dev-dependencies] @@ -995,12 +1010,14 @@ requires-dist = [ { name = "click" }, { name = "numpy", specifier = ">=2" }, { name = "pandas" }, + { name = "tabulate", specifier = ">=0.10.0" }, { name = "tqdm" }, { name = "tsinfer-hmm", editable = "../tsinfer-hmm" }, { name = "tskit", specifier = ">=1.0.0" }, { name = "tskit-hmm", editable = "../tskit-hmm" }, { name = "tsls", editable = "../tsls" }, { name = "tszip" }, + { name = "zarr", specifier = ">=3.1" }, ] [package.metadata.requires-dev]