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/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/README.md b/benchmarking/README.md index 74abc70..bdf9656 100644 --- a/benchmarking/README.md +++ b/benchmarking/README.md @@ -363,10 +363,242 @@ 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).) + 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 all 6,404 samples takes about 4.5 CPU hours with either engine. + +### 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 +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. + +### 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 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). + +``` +uv run python paired.py +``` + +#### Current result + +Two full runs of the current build (`832fc5f`), tsls / tsinfer-hmm: + +| workload | queries | round 1 | round 2 | +|---|---|---|---| +| 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 | + 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 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 + 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. +- **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 + +- 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 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 + 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 + +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, 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). diff --git a/benchmarking/callgrind.py b/benchmarking/callgrind.py new file mode 100644 index 0000000..6d29f43 --- /dev/null +++ b/benchmarking/callgrind.py @@ -0,0 +1,226 @@ +""" +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 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 + +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", +} + +# (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), + ("ancestors_g1483", 10, None), + ("ancestors_g1484", 3, None), + ("ancestors_g1485", 5, None), + ("samples", 1, 100_000), +] + + +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"], + ) + + +def cut_reference(ts, num_sites): + """ + Return the tree sequence cut just before the specified site, so that it + has only the first num_sites sites. Node IDs are unchanged. + """ + cut = ts.sites_position[num_sites] + return ts.keep_intervals([[0, cut]], simplify=False).rtrim() + + +@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) +@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, num_sites, instrument): + """ + Match the first NUM_QUERIES queries in WORKLOAD_DIR with the specified + 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( + ["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) + 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))) + + +def read_total_ir(path): + with open(path) as f: + for line in f: + if line.startswith(("totals:", "summary:")): + return int(line.split()[1]) + raise ValueError(f"No totals in {path}") + + +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, num_sites, engine_name, out_file): + cmd = [ + "valgrind", + "--tool=callgrind", + "--instr-atstart=no", + "--collect-atstart=no", + f"--toggle-collect={FIND_PATH_FUNCTIONS[engine_name]}", + f"--callgrind-out-file={out_file}", + sys.executable, + __file__, + "run", + str(pathlib.Path(data_dir) / workload), + engine_name, + "-n", + str(num_queries), + "--instrument", + ] + 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:]}") + paths_file = out_file.with_suffix(".paths") + paths_file.write_text(result.stdout) + 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( + "-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("-j", "--jobs", type=int, default=3) +def measure(data_dir, workload_names, out_dir, baseline_dir, jobs): + """ + 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 = WORKLOADS + if len(workload_names) > 0: + workloads = [w for w in WORKLOADS if w[0] in workload_names] + tasks = [] + 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, num_sites, engine_name) + if engine_name == "tsinfer" and path.exists(): + continue + 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: + 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, num_sites in workloads: + counts = [] + paths = [] + for engine_name, directory in [("tsinfer", baseline_dir), ("tsls", out_dir)]: + 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 num_sites is None + else f"{num_queries} ({num_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}") + + +if __name__ == "__main__": + cli() diff --git a/benchmarking/engines.py b/benchmarking/engines.py index 99b3274..bc56739 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,20 @@ 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(), alleles=alleles) + + def make_matcher(self): + return tsls.Hmm( + self.reference, + self.recombination, + self.mismatch, + 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/benchmarking/paired.py b/benchmarking/paired.py new file mode 100644 index 0000000..3f4164e --- /dev/null +++ b/benchmarking/paired.py @@ -0,0 +1,109 @@ +""" +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), +] + +# 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): + """ + 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) +@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 = 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( + "| 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() diff --git a/experiments/compression/README.md b/experiments/compression/README.md index 7883ccf..d043324 100644 --- a/experiments/compression/README.md +++ b/experiments/compression/README.md @@ -11,7 +11,11 @@ uv run python capture.py --sizes 10000,30000,100000 # about 35 minutes uv run python report.py # results/report.md ``` -`capture.py` uses `AncestorMatcher.viterbi_matrix` and `forward_matrix`. It +`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 `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)`. @@ -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..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 @@ -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..6089084 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,13 +58,13 @@ 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: + 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, 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) @@ -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,283 @@ 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.Hmm( + 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.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.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)) + 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.Hmm( + 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/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..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.MatcherIndexes(ts, alleles=v.ALLELES) + mi = tsls.TreeSequence(ts, alleles=v.ALLELES) 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, 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) 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, 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") @@ -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..617e84f 100644 --- a/tsls/_tslsmodule.c +++ b/tsls/_tslsmodule.c @@ -23,14 +23,15 @@ static PyObject *TslsMatchImpossible; typedef struct { PyObject_HEAD - matcher_indexes_t *matcher_indexes; -} MatcherIndexes; + tsk_treeseq_t *tree_sequence; + tsls_index_t *index; +} 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 +47,182 @@ 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, + * 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 || self->index == 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->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); + 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)) { - 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) { + LightweightTableCollection *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; + 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; } - 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"); + 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; } - 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; + 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; + } } - 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); + if (LightweightTableCollection_check_state(tables) != 0) { + goto out; } - - 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_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: - Py_XDECREF(num_alleles_array); - Py_XDECREF(ancestral_state_array); - Py_XDECREF(derived_state_array); + PyMem_Free(alleles_data); + Py_XDECREF(alleles_seq); 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,10 +231,10 @@ 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", + static char *kwlist[] = { "tree_sequence", "recombination", "mismatch", "extended_checks", "likelihood_threshold", "weight_by_n", "pruning_mode", "pruning_parameter", "count_distinct", NULL }; - MatcherIndexes *matcher_indexes = NULL; + TreeSequence *tree_sequence = NULL; PyObject *recombination = NULL; PyObject *mismatch = NULL; PyArrayObject *recombination_array = NULL; @@ -283,20 +242,22 @@ AncestorMatcher_init(AncestorMatcher *self, PyObject *args, PyObject *kwds) npy_intp *shape; double likelihood_threshold = DBL_MIN; int flags = 0; + size_t num_sites; - self->ancestor_matcher = NULL; - self->matcher_indexes = NULL; + self->hmm = NULL; + self->tree_sequence = NULL; if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!OO|idiidi", kwlist, - &MatcherIndexesType, &matcher_indexes, &recombination, &mismatch, + &TreeSequenceType, &tree_sequence, &recombination, &mismatch, &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); recombination_array = (PyArrayObject *) PyArray_FromAny(recombination, PyArray_DescrFromType(NPY_FLOAT64), 1, 1, NPY_ARRAY_IN_ARRAY, NULL); @@ -304,7 +265,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 +276,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,15 +295,14 @@ 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->index, + 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; @@ -355,7 +315,7 @@ AncestorMatcher_init(AncestorMatcher *self, PyObject *args, PyObject *kwds) } 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 +331,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 +348,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 +363,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 +466,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 +499,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 +508,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 +534,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 +553,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 +571,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 +623,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 +671,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 +685,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 +710,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 +751,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 +759,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 +851,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 new file mode 100644 index 0000000..c130d7e --- /dev/null +++ b/tsls/lib/hmm.c @@ -0,0 +1,1899 @@ +/* + * The Li and Stephens HMM on a tree sequence. + * + * The states of the HMM at a given position are all the nodes in the tree + * at that position, that is, the nodes with a parent or at least one child. + * Likelihoods are stored on a subset of the nodes: a node with + * NULL_LIKELIHOOD has the same likelihood as its nearest ancestor with a + * value. After processing each site we remove the value from any node + * whose value equals the value it would inherit from its parent. Every root + * always has a value. + * + * The same machinery runs the Viterbi, forward and backward algorithms. + * At each site, the probability of staying on a node is + * p * (1 - rho + rho / n) (Viterbi) or p * (1 - rho) + rho / n * S (forward + * and backward, where S is the total over all states at the previous site), + * and the probability of switching to a given node is rho / n. If weight_by_n is set n is the + * number of nodes in the tree sequence, otherwise n = 1. Values are + * normalised at each site (by the maximum for Viterbi and the sum for + * forward), and then values below likelihood_threshold are set to it. + * + * When a node enters the tree it gets the value 0, and so it can only be + * reached by recombination. When a node leaves the tree, the probability + * of continuing from it is only the probability of recombination. So a node + * can only be copied from one site to the next without recombination if it + * 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 +#include +#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 + * check_recombination_required argument. */ +#ifdef __GNUC__ +#define ALWAYS_INLINE inline __attribute__((always_inline)) +#else +#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 + +static inline void +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) +{ + const tsk_id_t lsib = left_sib[c]; + const tsk_id_t rsib = right_sib[c]; + + if (lsib == NULL_NODE) { + left_child[p] = rsib; + } else { + right_sib[lsib] = rsib; + } + if (rsib == NULL_NODE) { + right_child[p] = lsib; + } else { + left_sib[rsib] = lsib; + } + 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, 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]; + + tree[c].parent = p; + if (u == NULL_NODE) { + left_child[p] = c; + left_sib[c] = NULL_NODE; + right_sib[c] = NULL_NODE; + } else { + right_sib[u] = c; + left_sib[c] = u; + right_sib[c] = NULL_NODE; + } + right_child[p] = c; +} + +/* 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) + +/* 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 tsls_node_t *restrict tree, tsk_id_t u, allele_t ancestral_state) +{ + allele_t state; + + while (tree[u].allelic_state == TSK_NULL) { + u = tree[u].parent; + } + state = tree[u].allelic_state; + return likely(state != ROOT_ALLELE) ? state : ancestral_state; +} + +/* 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 +seek_tree_position(tsk_tree_position_t *tp, double x) +{ + int ret = 0; + 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 { + ret = tsk_tree_position_seek_backward(tp, index); + } + return ret; +} + +/*=================================================================== + * Compressed matrix + *=================================================================== + */ + +int +compressed_matrix_alloc(compressed_matrix_t *self, size_t num_sites, size_t block_size) +{ + int ret = 0; + + memset(self, 0, sizeof(*self)); + self->num_sites = num_sites; + self->normalisation_factor = calloc(num_sites + 1, sizeof(double)); + self->num_values = calloc(num_sites + 1, sizeof(tsk_size_t)); + self->nodes = calloc(num_sites + 1, sizeof(tsk_id_t *)); + self->values = calloc(num_sites + 1, sizeof(double *)); + self->flags = calloc(num_sites + 1, sizeof(int8_t *)); + if (self->normalisation_factor == NULL || self->num_values == NULL + || self->nodes == NULL || self->values == NULL || self->flags == NULL) { + ret = TSLS_ERR_NO_MEMORY; + goto out; + } + if (tsi_blkalloc_init(&self->allocator, TSK_MAX(block_size, 1024)) != 0) { + ret = TSLS_ERR_NO_MEMORY; + goto out; + } +out: + return ret; +} + +int +compressed_matrix_free(compressed_matrix_t *self) +{ + tsi_safe_free(self->normalisation_factor); + tsi_safe_free(self->num_values); + tsi_safe_free(self->nodes); + tsi_safe_free(self->values); + tsi_safe_free(self->flags); + tsi_blkalloc_free(&self->allocator); + return 0; +} + +size_t +compressed_matrix_get_total_values(const compressed_matrix_t *self) +{ + size_t j; + size_t total = 0; + + for (j = 0; j < self->num_sites; j++) { + total += self->num_values[j]; + } + return total; +} + +/* Round up to a multiple of 8 bytes so that the doubles stay aligned. */ +static inline size_t +aligned_size(size_t size) +{ + return (size + 7) & ~((size_t) 7); +} + +/* 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 tsls_likelihood_t *restrict L, const tsls_node_t *restrict tree, + bool with_flags) +{ + int ret = 0; + int j; + tsk_id_t *restrict nodes_out; + double *restrict values_out; + int8_t *restrict flags_out = NULL; + + if (with_flags) { + flags_out = tsi_blkalloc_get( + &self->allocator, aligned_size((size_t) num_values * sizeof(int8_t))); + if (flags_out == NULL) { + ret = TSLS_ERR_NO_MEMORY; + goto out; + } + for (j = 0; j < num_values; j++) { + flags_out[j] = tree[nodes[j]].recombination_required; + } + } + + values_out = tsi_blkalloc_get( + &self->allocator, aligned_size((size_t) num_values * sizeof(double))); + nodes_out = tsi_blkalloc_get( + &self->allocator, aligned_size((size_t) num_values * sizeof(tsk_id_t))); + if (values_out == NULL || nodes_out == NULL) { + ret = TSLS_ERR_NO_MEMORY; + goto out; + } + for (j = 0; j < num_values; j++) { + nodes_out[j] = nodes[j]; + values_out[j] = L[nodes[j]].value; + } + self->normalisation_factor[site] = normalisation_factor; + self->num_values[site] = (tsk_size_t) num_values; + self->nodes[site] = nodes_out; + self->values[site] = values_out; + self->flags[site] = flags_out; +out: + return ret; +} + +/*=================================================================== + * Index + *=================================================================== + */ + +/* 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_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; + 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; + tsk_size_t mutation_offset = 0; + 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++) { + /* 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, + 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; + } + 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; +} + +/* 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) +{ + int ret = 0; + const size_t num_mutations = tsk_treeseq_get_num_mutations(tree_sequence); + + 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_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->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)); + 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 + || self->removal_edges == NULL) { + ret = TSLS_ERR_NO_MEMORY; + goto out; + } + ret = tsls_index_encode_alleles(self, alleles, num_alleles); + if (ret != 0) { + goto out; + } + tsls_index_copy_edges(self); +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); + tsls_large_free_array(self->insertion_edges, self->num_edges); + tsls_large_free_array(self->removal_edges, self->num_edges); + return 0; +} + +/*=================================================================== + * HMM + *=================================================================== + */ + +static inline bool +tsls_hmm_in_tree(const tsls_hmm_t *self, tsk_id_t u) +{ + return self->tree[u].parent != 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].value >= 0); + } + for (u = 0; u < (tsk_id_t) self->num_nodes; u++) { + if (self->likelihood[u].value >= 0) { + num_likelihoods++; + } + if (tsls_hmm_in_tree(self, u)) { + assert(self->likelihood[u].value != NOT_IN_TREE); + if (self->tree[u].parent == NULL_NODE) { + /* Roots always have a likelihood */ + assert(self->likelihood[u].value >= 0); + } + } else { + assert(self->likelihood[u].value == NOT_IN_TREE); + } + 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; +} + +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->tree[j].parent, + self->left_child[j], self->right_child[j], self->left_sib[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].value); + } + 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)); + 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->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_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->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->likelihood_nodes == NULL + || self->likelihood_nodes_tmp == NULL || self->values_tmp == NULL) { + ret = TSLS_ERR_NO_MEMORY; + 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 + = TSK_MAX(64 * 1024 * 1024, aligned_size((N + 1) * sizeof(tsk_id_t)) * 2); + /* If the traceback allocator is using more than 2GiB of RAM free it, so + * that other threads can use the memory */ + self->traceback_realloc_size = 2L * 1024L * 1024L * 1024L; + if (tsi_blkalloc_init(&self->traceback_allocator, self->traceback_block_size) + != 0) { + ret = TSLS_ERR_NO_MEMORY; + goto out; + } + memcpy(self->recombination_rate, recombination_rate, + self->num_sites * sizeof(*self->recombination_rate)); + memcpy(self->mismatch_rate, mismatch_rate, + self->num_sites * sizeof(*self->mismatch_rate)); + memset(self->max_likelihood_node, 0xff, self->num_sites * sizeof(tsk_id_t)); +out: + return ret; +} + +int +tsls_hmm_free(tsls_hmm_t *self) +{ + tsi_safe_free(self->recombination_rate); + tsi_safe_free(self->mismatch_rate); + 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->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); + tsi_safe_free(self->max_likelihood_node); + tsi_safe_free(self->traceback); + tsi_blkalloc_free(&self->traceback_allocator); + return 0; +} + +int +tsls_hmm_set_pruning(tsls_hmm_t *self, int mode, double parameter) +{ + int ret = 0; + + if (mode == TSLS_PRUNE_NONE) { + parameter = 0; + } else if (mode == TSLS_PRUNE_RELATIVE) { + if (!(parameter >= 0 && parameter <= 1)) { + ret = TSLS_ERR_BAD_PRUNING; + goto out; + } + } else if (mode == TSLS_PRUNE_TOP_N) { + if (!(parameter >= 1)) { + ret = TSLS_ERR_BAD_PRUNING; + goto out; + } + } else { + ret = TSLS_ERR_BAD_PRUNING; + goto out; + } + self->pruning_mode = mode; + self->pruning_parameter = parameter; +out: + return ret; +} + +static int +cmp_double(const void *a, const void *b) +{ + const double x = *(const double *) a; + const double y = *(const double *) b; + return (x > y) - (x < y); +} + +/* Copy the current likelihood values into values_tmp, sorted. */ +static void +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]].value; + } + qsort(self->values_tmp, (size_t) self->num_likelihood_nodes, sizeof(double), + cmp_double); +} + +/* Apply the experimental pruning, giving pruned nodes the value + * likelihood_threshold. The values must be normalised. */ +static void +tsls_hmm_prune(tsls_hmm_t *self) +{ + 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; + int j, n; + + if (self->pruning_mode == TSLS_PRUNE_NONE || num_likelihood_nodes == 0) { + return; + } + 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]].value); + } + cutoff = self->pruning_parameter * max_L; + } else { + n = (int) self->pruning_parameter; + if (num_likelihood_nodes <= n) { + return; + } + tsls_hmm_sort_values(self); + /* Keep values >= the nth largest (ties at the cutoff are kept) */ + cutoff = self->values_tmp[num_likelihood_nodes - n]; + } + for (j = 0; j < num_likelihood_nodes; j++) { + if (L[L_nodes[j]].value < cutoff) { + L[L_nodes[j]].value = self->likelihood_threshold; + } + } +} + +/* Count the number of distinct likelihood values at this site. */ +static void +tsls_hmm_count_distinct(tsls_hmm_t *self) +{ + int j; + size_t num_distinct; + + if (!(self->flags & TSLS_COUNT_DISTINCT) || self->num_likelihood_nodes == 0) { + return; + } + 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]) { + num_distinct++; + } + } + self->total_distinct_values += num_distinct; +} + +static inline void +tsls_hmm_add_likelihood(tsls_hmm_t *self, tsk_id_t u, double value) +{ + self->likelihood[u].value = value; + self->likelihood_nodes[self->num_likelihood_nodes] = u; + self->num_likelihood_nodes++; +} + +/* 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_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; + + k = 0; + for (j = 0; j < self->num_likelihood_nodes; j++) { + L_nodes[k] = L_nodes[j]; + if (L[L_nodes[j]].value != NOT_IN_TREE) { + k++; + } + } + self->num_likelihood_nodes = k; +} + +static inline void +tsls_hmm_update_num_below(tsls_hmm_t *self, tsk_id_t u, int64_t delta) +{ + tsls_node_t *restrict tree = self->tree; + int64_t *restrict num_below = self->num_below; + + for (; u != NULL_NODE; u = tree[u].parent) { + num_below[u] += delta; + } +} + +/* Returns the likelihood inherited by node u, filling in the likelihood + * 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 inline double +tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u) +{ + 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].value == NULL_LIKELIHOOD) && likely(L[v].cache == CACHE_UNSET)) { + v = tree[v].parent; + } + value = L[v].cache; + if (unlikely(value == CACHE_UNSET)) { + value = L[v].value; + } + v = u; + while (likely(L[v].value == NULL_LIKELIHOOD) && likely(L[v].cache == CACHE_UNSET)) { + L[v].cache = value; + v = tree[v].parent; + } + 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) +{ + tsls_likelihood_t *restrict L = self->likelihood; + tsls_node_t *restrict tree = self->tree; + + while (L[u].cache != CACHE_UNSET) { + L[u].cache = CACHE_UNSET; + u = tree[u].parent; + } +} + +/* 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 tsls_hmm_update_tree). */ +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) +{ + 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].value == NULL_LIKELIHOOD) { + tsls_hmm_add_likelihood(self, c, tsls_hmm_get_inherited_likelihood(self, p)); + } + remove_edge(p, c, self->tree, self->left_child, self->right_child, + self->left_sib, self->right_sib); + 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[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. */ + if (!tsls_hmm_in_tree(self, c)) { + isolated[*num_isolated] = c; + (*num_isolated)++; + } + if (!tsls_hmm_in_tree(self, p)) { + /* p was a root, and so has a likelihood */ + assert(L[p].value >= 0); + isolated[*num_isolated] = p; + (*num_isolated)++; + } + } +} + +/* 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 inline void +tsls_hmm_insert_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, bool update_likelihoods) +{ + tsls_likelihood_t *restrict L = self->likelihood; + + if (update_likelihoods) { + if (L[p].value == NOT_IN_TREE) { + tsls_hmm_add_likelihood(self, p, 0); + } + 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].value >= 0); + } + insert_edge(p, c, self->tree, self->left_child, self->right_child, + self->left_sib, self->right_sib); + self->tree[c].allelic_state = TSK_NULL; + if (self->track_num_below) { + tsls_hmm_update_num_below(self, p, self->num_below[c]); + } +} + +/* 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 +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; + tsk_id_t j, u; + int k, num_left, num_isolated = 0; + + for (j = tp->out.start; j != tp->out.stop; j += direction) { + tsls_hmm_remove_edge( + 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, out[j].parent); + } + for (j = tp->in.start; j != tp->in.stop; j += direction) { + 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 + * 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)) { + L[u].value = 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, + * 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; + const tsls_edge_t *restrict in; + tsk_id_t u, j; + size_t k; + + 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)); + for (k = 0; k < N; k++) { + /* 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); + 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) { + 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; + for (u = 0; u < (tsk_id_t) N; u++) { + 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].value = NOT_IN_TREE; + } + } +out: + return ret; +} + +static inline void +tsls_hmm_set_allelic_state(tsls_hmm_t *self, const tsk_id_t 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 = index->site_mutations_offset[site]; j < stop; j++) { + self->tree[node[j]].allelic_state = derived_state[j]; + } +} + +/* 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 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 = index->site_mutations_offset[site]; j < stop; j++) { + u = node[j]; + self->tree[u].allelic_state = self->tree[u].parent == 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_clear_allelic_state(tsls_hmm_t *self, const tsk_id_t 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 = index->site_mutations_offset[site]; j < stop; j++) { + self->tree[node[j]].allelic_state = TSK_NULL; + } +} + +/* Make sure that every node with a mutation at this site that is in the + * tree has its own likelihood. */ +static void +tsls_hmm_insert_mutation_likelihoods(tsls_hmm_t *self, const tsk_id_t site) +{ + 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; + tsk_id_t u, node; + + for (j = index->site_mutations_offset[site]; j < stop; j++) { + node = index->mutation_node[j]; + if (L[node].value == NULL_LIKELIHOOD) { + u = node; + while (L[u].value == NULL_LIKELIHOOD) { + u = tree[u].parent; + } + tsls_hmm_add_likelihood(self, node, L[u].value); + } + } +} + +/* 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 ALWAYS_INLINE double +tsls_hmm_get_inherited_likelihood_and_flag( + tsls_hmm_t *self, tsk_id_t u, int8_t *flag) +{ + 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].value == NULL_LIKELIHOOD) && likely(L[v].cache == CACHE_UNSET)) { + v = tree[v].parent; + } + if (L[v].cache == CACHE_UNSET) { + value = L[v].value; + r = tree[v].recombination_required; + } else { + value = L[v].cache; + r = tree[v].recombination_required_cache; + } + v = u; + 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; +} + +/* Remove likelihoods that are equal to the likelihood inherited from the + * 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 ALWAYS_INLINE void +tsls_hmm_coalesce_likelihoods( + tsls_hmm_t *self, bool check_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; + 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; + int j, num_cached_paths, num_likelihood_nodes; + + num_cached_paths = 0; + num_likelihood_nodes = 0; + for (j = 0; j < old_num_likelihood_nodes; j++) { + u = L_nodes[j]; + p = tree[u].parent; + if (p != NULL_NODE) { + cached_paths[num_cached_paths] = p; + num_cached_paths++; + if (check_recombination_required) { + /* A removed node had the inherited value and flag, so the + * cached values below it stay valid. */ + L_p = tsls_hmm_get_inherited_likelihood_and_flag( + self, p, &R_p); + 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].value == L_p) { + L[u].value = NULL_LIKELIHOOD; + } + } + } + if (L[u].value >= 0) { + L_nodes[num_likelihood_nodes] = L_nodes[j]; + num_likelihood_nodes++; + } + } + self->num_likelihood_nodes = num_likelihood_nodes; + for (j = 0; j < num_cached_paths; j++) { + tsls_hmm_reset_likelihood_cache(self, cached_paths[j]); + } +} + +static inline double +emission_probability(allele_t node_state, allele_t state, double mu, double num_alleles) +{ + if (state == TSLS_MISSING_DATA || node_state == state) { + return 1 - (num_alleles - 1) * mu; + } + return mu; +} + +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->index->num_alleles[site]; + + if (state < TSLS_MISSING_DATA || (state >= 0 && (tsk_size_t) state >= num_alleles)) { + return TSLS_ERR_BAD_HAPLOTYPE_ALLELE; + } + return 0; +} + +static inline double +tsls_hmm_get_n(const tsls_hmm_t *self) +{ + double n = 1.0; + if (!(self->flags & TSLS_DISABLE_WEIGHT_BY_N)) { + n = (double) self->num_nodes; + } + return n; +} + +/* Compute the number of nodes in the tree that inherit the likelihood of + * each likelihood node. */ +static void +tsls_hmm_compute_num_inheriting(tsls_hmm_t *self) +{ + 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; + tsk_id_t u, w; + int j; + + for (j = 0; j < self->num_likelihood_nodes; j++) { + u = L_nodes[j]; + num_inheriting[u] = num_below[u]; + } + for (j = 0; j < self->num_likelihood_nodes; j++) { + u = L_nodes[j]; + w = tree[u].parent; + if (w != NULL_NODE) { + while (L[w].value == NULL_LIKELIHOOD) { + w = tree[w].parent; + } + num_inheriting[w] -= num_below[u]; + } + } +} + +/* Returns the sum of the likelihoods over all the nodes in the tree. */ +static double +tsls_hmm_total_likelihood(const tsls_hmm_t *self) +{ + 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; + tsk_id_t u; + int j; + + 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].value; + } + return total; +} + +/* Store the recombination_required state in the traceback */ +static int WARN_UNUSED +tsls_hmm_store_traceback(tsls_hmm_t *self, const tsk_id_t site_id) +{ + int ret = 0; + tsk_id_t u; + int j; + int8_t *restrict list_R; + tsk_id_t *restrict list_node; + node_state_list_t *restrict list; + node_state_list_t *restrict T = self->traceback; + const tsk_id_t *restrict nodes = self->likelihood_nodes; + tsls_node_t *restrict tree = self->tree; + const int num_likelihood_nodes = self->num_likelihood_nodes; + bool match; + + /* Check to see if the previous site has the same recombination_required. If so, + * we can reuse the same list. */ + match = false; + if (site_id > 0) { + list = &T[site_id - 1]; + if (list->size == num_likelihood_nodes) { + list_node = list->node; + list_R = list->recombination_required; + match = true; + for (j = 0; j < num_likelihood_nodes; j++) { + if (list_node[j] != nodes[j] || list_R[j] != tree[nodes[j]].recombination_required) { + match = false; + break; + } + } + } + } + + if (match) { + T[site_id].size = T[site_id - 1].size; + T[site_id].node = T[site_id - 1].node; + T[site_id].recombination_required = T[site_id - 1].recombination_required; + } else { + list_node = tsi_blkalloc_get(&self->traceback_allocator, + aligned_size((size_t) num_likelihood_nodes * sizeof(tsk_id_t))); + list_R = tsi_blkalloc_get(&self->traceback_allocator, + aligned_size((size_t) num_likelihood_nodes * sizeof(int8_t))); + if (list_node == NULL || list_R == NULL) { + ret = TSLS_ERR_NO_MEMORY; + goto out; + } + T[site_id].node = list_node; + T[site_id].recombination_required = list_R; + T[site_id].size = num_likelihood_nodes; + for (j = 0; j < num_likelihood_nodes; j++) { + u = nodes[j]; + list_node[j] = u; + list_R[j] = tree[u].recombination_required; + } + } + self->total_traceback_size += (size_t) num_likelihood_nodes; +out: + return ret; +} + +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; + 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); + 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; + + 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++) { + u = L_nodes[j]; + p_no_recomb = L[u].value * (1 - rho + rho / n); + p_recomb = rho / n; + tree[u].recombination_required = false; + if (p_no_recomb > p_recomb) { + p_t = p_no_recomb; + } else { + p_t = p_recomb; + tree[u].recombination_required = true; + } + p_e = emission_probability( + get_marked_allelic_state(tree, u, ancestral_state), state, + mu, num_alleles); + L[u].value = p_t * p_e; + if (L[u].value > max_L) { + max_L = L[u].value; + max_L_node = u; + } + } + tsls_hmm_unset_allelic_state(self, site); + (void) num_likelihood_nodes; + + if (max_L <= 0) { + if (mu <= 0 || mu >= 1) { + ret = TSLS_ERR_MATCH_IMPOSSIBLE_EXTREME_MUTATION_PROBA; + } else if (rho == 0) { + ret = TSLS_ERR_MATCH_IMPOSSIBLE_ZERO_RECOMB_PRECISION; + } else { + ret = TSLS_ERR_MATCH_IMPOSSIBLE; + } + goto out; + } + assert(max_L_node != NULL_NODE); + self->max_likelihood_node[site] = max_L_node; + for (j = 0; j < self->num_likelihood_nodes; j++) { + u = L_nodes[j]; + L[u].value = TSK_MAX(L[u].value / max_L, self->likelihood_threshold); + } + 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. */ + tsls_hmm_coalesce_likelihoods(self, true); + ret = tsls_hmm_store_traceback(self, site); + if (ret != 0) { + goto out; + } + if (output != NULL) { + ret = compressed_matrix_store_site( + output, site, max_L, self->num_likelihood_nodes, L_nodes, L, self->tree, true); + if (ret != 0) { + goto out; + } + } +out: + return ret; +} + +/* Recombination redistributes the total probability over the states at the + * previous site, including nodes that have since left the tree, so we + * 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 +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; + 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); + 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; + + tsls_hmm_insert_mutation_likelihoods(self, site); + tsls_hmm_compute_num_inheriting(self); + total = *previous_total; + if (total < 0) { + total = tsls_hmm_total_likelihood(self); + } + 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(tree, u, ancestral_state), state, + mu, num_alleles); + 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); + if (!(normalisation_factor > 0)) { + ret = TSLS_ERR_MATCH_IMPOSSIBLE; + goto out; + } + for (j = 0; j < self->num_likelihood_nodes; j++) { + u = L_nodes[j]; + 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, self->tree, false); + if (ret != 0) { + goto out; + } + tsls_hmm_coalesce_likelihoods(self, false); +out: + return ret; +} + +/* The backward transition from site l to site l - 1 depends on the tree at + * site l - 1 (nodes may have entered or left the tree), so we apply it + * lazily when we reach site l - 1, using these saved values from site l. + */ +typedef struct { + bool pending; + double total; + double rho; + double normalisation_factor; +} backward_transition_t; + +static int WARN_UNUSED +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; + 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]; + const allele_t ancestral_state = self->index->ancestral_state[site]; + double x; + tsk_id_t u; + int j; + + if (transition->pending) { + for (j = 0; j < self->num_likelihood_nodes; j++) { + u = L_nodes[j]; + 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, self->tree, false); + if (ret != 0) { + goto out; + } + 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].value *= emission_probability( + get_marked_allelic_state(tree, u, ancestral_state), state, + mu, num_alleles); + } + tsls_hmm_unset_allelic_state(self, site); + transition->pending = true; + transition->total = tsls_hmm_total_likelihood(self); + transition->rho = self->recombination_rate[site]; + transition->normalisation_factor = forward_normalisation_factor[site]; + tsls_hmm_coalesce_likelihoods(self, false); +out: + return ret; +} + +/* Run the specified algorithm over sites [start, end). */ +static int WARN_UNUSED +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; + tsk_tree_position_t tp; + backward_transition_t transition; + 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) { + goto out; + } + } + self->track_num_below = mode != MODE_VITERBI; + first_site = direction > 0 ? start : end - 1; + 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.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) { + tsls_hmm_check_state(self); + } + if (mode == MODE_VITERBI) { + ret = tsls_hmm_process_site_viterbi(self, site, haplotype[site], output); + } else if (mode == MODE_FORWARD) { + ret = tsls_hmm_process_site_forward( + self, site, haplotype[site], &previous_total, output); + } else { + ret = tsls_hmm_process_site_backward(self, site, haplotype[site], + forward_normalisation_factor, &transition, output); + } + if (ret != 0) { + goto out; + } + } +out: + tsk_tree_position_free(&tp); + return ret; +} + +/* Resets the recombination_required array from the traceback at the specified site. + */ +static inline void +tsls_hmm_set_recombination_required( + tsls_hmm_t *self, tsk_id_t site, tsls_node_t *restrict tree) +{ + int j; + 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++) { + tree[list_node[j]].recombination_required = list_R[j]; + } +} + +/* Unsets the likelihood array from the traceback at the specified site. + */ +static inline void +tsls_hmm_unset_recombination_required( + tsls_hmm_t *self, tsk_id_t site, tsls_node_t *restrict tree) +{ + int j; + const tsk_id_t *restrict list_node = self->traceback[site].node; + const int size = self->traceback[site].size; + + for (j = 0; j < size; j++) { + tree[list_node[j]].recombination_required = -1; + } +} + +/* 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) +{ + int ret = 0; + tsk_tree_position_t tp; + tsk_id_t site, j, u, max_likelihood_node; + double x; + bool valid; + 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 + * 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 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]; + size_t path_length = 0; + + tsk_tree_position_init(&tp, self->tree_sequence, 0); + + /* Prepare for the traceback and get the memory ready for recording + * the output edges. */ + path_right[path_length] = end_pos; + max_likelihood_node = self->max_likelihood_node[end - 1]; + assert(max_likelihood_node != NULL_NODE); + path_parent[path_length] = max_likelihood_node; + + /* 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++) { + tree[u].parent = virtual_root; + tree[u].allelic_state = TSK_NULL; + tree[u].recombination_required = -1; + } + tree[virtual_root].parent = NULL_NODE; + x = sites_position[end - 1]; + ret = seek_tree_position(&tp, x); + 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) { + if (seek_edge_covers(&tp, &in[j], x)) { + tree[in[j].child].parent = in[j].parent; + } + } + for (site = end - 1; site >= start; site--) { + while (sites_position[site] < tp.interval.left) { + 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--) { + tree[out[j].child].parent = virtual_root; + } + for (j = tp.in.start; j != tp.in.stop; j--) { + tree[in[j].child].parent = in[j].parent; + } + } + tree[virtual_root].allelic_state = self->index->ancestral_state[site]; + tsls_hmm_set_allelic_state(self, site); + u = path_parent[path_length]; + while (tree[u].allelic_state == TSK_NULL) { + u = tree[u].parent; + } + 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, 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 (tree[u].recombination_required == -1) { + u = tree[u].parent; + assert(u != virtual_root); + } + 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]; + 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, tree); + } + path_left[path_length] = start_pos; + path_length++; + *path_length_out = path_length; +out: + tsk_tree_position_free(&tp); + return ret; +} + +static int +tsls_hmm_reset_traceback(tsls_hmm_t *self) +{ + int ret = 0; + + if (self->traceback_allocator.total_size > self->traceback_realloc_size) { + tsi_blkalloc_free(&self->traceback_allocator); + if (tsi_blkalloc_init(&self->traceback_allocator, self->traceback_block_size) + != 0) { + ret = TSLS_ERR_NO_MEMORY; + goto out; + } + } else { + tsi_blkalloc_reset(&self->traceback_allocator); + } + self->total_traceback_size = 0; + self->total_distinct_values = 0; +out: + return ret; +} + +/* Run the Viterbi algorithm and traceback. If output is not NULL, the + * normalised values and recombination_required flags stored in the + * traceback at each site are also stored in output, with the maximum + * likelihood (before normalisation) as the normalisation factor. */ +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, double *path_left, double *path_right, + 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; + goto out; + } + if (output != NULL && output->num_sites != self->num_sites) { + ret = TSLS_ERR_BAD_SITE_RANGE; + goto out; + } + valid_window = true; + ret = tsls_hmm_reset_traceback(self); + if (ret != 0) { + goto out; + } + ret = tsls_hmm_run(self, MODE_VITERBI, start, end, haplotype, NULL, output); + if (ret != 0) { + goto out; + } + ret = tsls_hmm_run_traceback(self, start, end, matched_haplotype, path_length, + path_left, path_right, path_parent); +out: + /* 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)); + } + return ret; +} + +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) +{ + return tsls_hmm_find_path_matrix(self, start, end, haplotype, + matched_haplotype, path_length, path_left, path_right, path_parent, NULL); +} + +int +tsls_hmm_forward( + tsls_hmm_t *self, const allele_t *haplotype, compressed_matrix_t *output) +{ + int ret = 0; + + if (output->num_sites != self->num_sites) { + ret = TSLS_ERR_BAD_SITE_RANGE; + goto out; + } + if (self->num_sites > 0) { + ret = tsls_hmm_run(self, MODE_FORWARD, 0, (tsk_id_t) self->num_sites, + haplotype, NULL, output); + } +out: + return ret; +} + +int +tsls_hmm_backward(tsls_hmm_t *self, const allele_t *haplotype, + const double *forward_normalisation_factor, compressed_matrix_t *output) +{ + int ret = 0; + size_t j; + + if (output->num_sites != self->num_sites) { + ret = TSLS_ERR_BAD_SITE_RANGE; + goto out; + } + for (j = 0; j < self->num_sites; j++) { + if (!(forward_normalisation_factor[j] > 0)) { + ret = TSLS_ERR_BAD_NORMALISATION_FACTOR; + goto out; + } + } + if (self->num_sites > 0) { + ret = tsls_hmm_run(self, MODE_BACKWARD, 0, (tsk_id_t) self->num_sites, + haplotype, forward_normalisation_factor, output); + } +out: + return ret; +} + +double +tsls_hmm_get_mean_traceback_size(tsls_hmm_t *self) +{ + if (self->num_sites == 0) { + return 0; + } + return (double) self->total_traceback_size / ((double) self->num_sites); +} + +double +tsls_hmm_get_mean_num_distinct_values(tsls_hmm_t *self) +{ + if (self->num_sites == 0) { + return 0; + } + return (double) self->total_distinct_values / ((double) self->num_sites); +} + +size_t +tsls_hmm_get_total_memory(tsls_hmm_t *self) +{ + return self->traceback_allocator.total_size; +} diff --git a/tsls/lib/matcher.c b/tsls/lib/matcher.c deleted file mode 100644 index 1eb9307..0000000 --- a/tsls/lib/matcher.c +++ /dev/null @@ -1,1763 +0,0 @@ -/* - * The Li and Stephens HMM on a tree sequence. - * - * The states of the HMM at a given position are all the nodes in the tree - * at that position, that is, the nodes with a parent or at least one child. - * Likelihoods are stored on a subset of the nodes: a node with - * NULL_LIKELIHOOD has the same likelihood as its nearest ancestor with a - * value. After processing each site we remove the value from any node - * whose value equals the value it would inherit from its parent. Every root - * always has a value. - * - * The same machinery runs the Viterbi, forward and backward algorithms. - * At each site, the probability of staying on a node is - * p * (1 - rho + rho / n) (Viterbi) or p * (1 - rho) + rho / n * S (forward - * and backward, where S is the total over all states at the previous site), - * and the probability of switching to a given node is rho / n. If weight_by_n is set n is the - * number of nodes in the tree sequence, otherwise n = 1. Values are - * normalised at each site (by the maximum for Viterbi and the sum for - * forward), and then values below likelihood_threshold are set to it. - * - * When a node enters the tree it gets the value 0, and so it can only be - * reached by recombination. When a node leaves the tree, the probability - * of continuing from it is only the probability of recombination. So a node - * can only be copied from one site to the next without recombination if it - * is in the tree over the whole interval between them. - */ - -#include "tsls.h" - -#include -#include -#include -#include -#include - -#define MODE_VITERBI 0 -#define MODE_FORWARD 1 -#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) -{ - 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]; - - if (lsib == NULL_NODE) { - left_child[p] = rsib; - } else { - right_sib[lsib] = rsib; - } - if (rsib == NULL_NODE) { - right_child[p] = lsib; - } else { - left_sib[rsib] = lsib; - } - parent[c] = NULL_NODE; - left_sib[c] = NULL_NODE; - right_sib[c] = NULL_NODE; -} - -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) -{ - 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; - if (u == NULL_NODE) { - left_child[p] = c; - left_sib[c] = NULL_NODE; - right_sib[c] = NULL_NODE; - } else { - right_sib[u] = c; - left_sib[c] = u; - right_sib[c] = NULL_NODE; - } - 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]; -} - -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; - } -} - -static int WARN_UNUSED -matcher_indexes_add_mutation( - matcher_indexes_t *self, tsk_id_t site, tsk_id_t node, allele_t derived_state) -{ - 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; - } else { - tail = self->sites.mutations[site]; - while (tail->next != NULL) { - tail = tail->next; - } - tail->next = list_node; - } - 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 - *=================================================================== - */ - -int -compressed_matrix_alloc(compressed_matrix_t *self, size_t num_sites, size_t block_size) -{ - int ret = 0; - - memset(self, 0, sizeof(*self)); - self->num_sites = num_sites; - self->normalisation_factor = calloc(num_sites + 1, sizeof(double)); - self->num_values = calloc(num_sites + 1, sizeof(tsk_size_t)); - self->nodes = calloc(num_sites + 1, sizeof(tsk_id_t *)); - self->values = calloc(num_sites + 1, sizeof(double *)); - self->flags = calloc(num_sites + 1, sizeof(int8_t *)); - if (self->normalisation_factor == NULL || self->num_values == NULL - || self->nodes == NULL || self->values == NULL || self->flags == NULL) { - ret = TSLS_ERR_NO_MEMORY; - goto out; - } - if (tsi_blkalloc_init(&self->allocator, TSK_MAX(block_size, 1024)) != 0) { - ret = TSLS_ERR_NO_MEMORY; - goto out; - } -out: - return ret; -} - -int -compressed_matrix_free(compressed_matrix_t *self) -{ - tsi_safe_free(self->normalisation_factor); - tsi_safe_free(self->num_values); - tsi_safe_free(self->nodes); - tsi_safe_free(self->values); - tsi_safe_free(self->flags); - tsi_blkalloc_free(&self->allocator); - return 0; -} - -size_t -compressed_matrix_get_total_values(const compressed_matrix_t *self) -{ - size_t j; - size_t total = 0; - - for (j = 0; j < self->num_sites; j++) { - total += self->num_values[j]; - } - return total; -} - -/* Round up to a multiple of 8 bytes so that the doubles stay aligned. */ -static inline size_t -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. */ -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) -{ - int ret = 0; - int j; - tsk_id_t *restrict nodes_out; - double *restrict values_out; - int8_t *restrict flags_out = NULL; - - if (R != NULL) { - flags_out = tsi_blkalloc_get( - &self->allocator, aligned_size((size_t) num_values * sizeof(int8_t))); - if (flags_out == NULL) { - ret = TSLS_ERR_NO_MEMORY; - goto out; - } - for (j = 0; j < num_values; j++) { - flags_out[j] = R[nodes[j]]; - } - } - - values_out = tsi_blkalloc_get( - &self->allocator, aligned_size((size_t) num_values * sizeof(double))); - nodes_out = tsi_blkalloc_get( - &self->allocator, aligned_size((size_t) num_values * sizeof(tsk_id_t))); - if (values_out == NULL || nodes_out == NULL) { - ret = TSLS_ERR_NO_MEMORY; - goto out; - } - for (j = 0; j < num_values; j++) { - nodes_out[j] = nodes[j]; - values_out[j] = L[nodes[j]]; - } - self->normalisation_factor[site] = normalisation_factor; - self->num_values[site] = (tsk_size_t) num_values; - self->nodes[site] = nodes_out; - self->values[site] = values_out; - self->flags[site] = flags_out; -out: - return ret; -} - -/*=================================================================== - * Ancestor matcher - *=================================================================== - */ - -static inline bool -ancestor_matcher_in_tree(const ancestor_matcher_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) -{ - 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 (ancestor_matcher_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] == TSK_NULL); - assert(self->likelihood_cache[u] == CACHE_UNSET); - } - assert(num_likelihoods == self->num_likelihood_nodes); - (void) num_likelihoods; -} - -int -ancestor_matcher_print_state(ancestor_matcher_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"); - 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, "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 -ancestor_matcher_alloc(ancestor_matcher_t *self, - const matcher_indexes_t *matcher_indexes, const double *recombination_rate, - const double *mismatch_rate, double likelihood_threshold, int flags) -{ - int ret = 0; - const size_t N = matcher_indexes->num_nodes; - - memset(self, 0, sizeof(ancestor_matcher_t)); - self->flags = flags; - self->likelihood_threshold = likelihood_threshold; - self->matcher_indexes = matcher_indexes; - self->num_sites = matcher_indexes->num_sites; - self->num_nodes = N; - self->recombination_rate - = malloc((self->num_sites + 1) * sizeof(*self->recombination_rate)); - self->mismatch_rate = malloc((self->num_sites + 1) * sizeof(*self->mismatch_rate)); - 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->recombination_required - = malloc((N + 1) * sizeof(*self->recombination_required)); - self->likelihood_cache = malloc((N + 1) * sizeof(*self->likelihood_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)); - - 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 - || 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->likelihood_nodes == NULL || self->likelihood_nodes_tmp == NULL - || self->values_tmp == NULL || self->allelic_state == NULL) { - ret = TSLS_ERR_NO_MEMORY; - 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 - = TSK_MAX(64 * 1024 * 1024, aligned_size((N + 1) * sizeof(tsk_id_t)) * 2); - /* If the traceback allocator is using more than 2GiB of RAM free it, so - * that other threads can use the memory */ - self->traceback_realloc_size = 2L * 1024L * 1024L * 1024L; - if (tsi_blkalloc_init(&self->traceback_allocator, self->traceback_block_size) - != 0) { - ret = TSLS_ERR_NO_MEMORY; - goto out; - } - memcpy(self->recombination_rate, recombination_rate, - 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; -} - -int -ancestor_matcher_free(ancestor_matcher_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->likelihood_nodes); - tsi_safe_free(self->likelihood_nodes_tmp); - tsi_safe_free(self->values_tmp); - tsi_safe_free(self->allelic_state); - tsi_safe_free(self->max_likelihood_node); - tsi_safe_free(self->traceback); - tsi_blkalloc_free(&self->traceback_allocator); - return 0; -} - -int -ancestor_matcher_set_pruning(ancestor_matcher_t *self, int mode, double parameter) -{ - int ret = 0; - - if (mode == TSLS_PRUNE_NONE) { - parameter = 0; - } else if (mode == TSLS_PRUNE_RELATIVE) { - if (!(parameter >= 0 && parameter <= 1)) { - ret = TSLS_ERR_BAD_PRUNING; - goto out; - } - } else if (mode == TSLS_PRUNE_TOP_N) { - if (!(parameter >= 1)) { - ret = TSLS_ERR_BAD_PRUNING; - goto out; - } - } else { - ret = TSLS_ERR_BAD_PRUNING; - goto out; - } - self->pruning_mode = mode; - self->pruning_parameter = parameter; -out: - return ret; -} - -static int -cmp_double(const void *a, const void *b) -{ - const double x = *(const double *) a; - const double y = *(const double *) b; - return (x > y) - (x < y); -} - -/* Copy the current likelihood values into values_tmp, sorted. */ -static void -ancestor_matcher_sort_values(ancestor_matcher_t *self) -{ - int j; - - for (j = 0; j < self->num_likelihood_nodes; j++) { - self->values_tmp[j] = self->likelihood[self->likelihood_nodes[j]]; - } - qsort(self->values_tmp, (size_t) self->num_likelihood_nodes, sizeof(double), - cmp_double); -} - -/* 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) -{ - double *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; - int j, n; - - if (self->pruning_mode == TSLS_PRUNE_NONE || num_likelihood_nodes == 0) { - return; - } - 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]]); - } - cutoff = self->pruning_parameter * max_L; - } else { - n = (int) self->pruning_parameter; - if (num_likelihood_nodes <= n) { - return; - } - ancestor_matcher_sort_values(self); - /* Keep values >= the nth largest (ties at the cutoff are kept) */ - 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; - } - } -} - -/* Count the number of distinct likelihood values at this site. */ -static void -ancestor_matcher_count_distinct(ancestor_matcher_t *self) -{ - int j; - size_t num_distinct; - - if (!(self->flags & TSLS_COUNT_DISTINCT) || self->num_likelihood_nodes == 0) { - return; - } - ancestor_matcher_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]) { - num_distinct++; - } - } - self->total_distinct_values += num_distinct; -} - -static inline void -ancestor_matcher_add_likelihood(ancestor_matcher_t *self, tsk_id_t u, double value) -{ - self->likelihood[u] = value; - self->likelihood_nodes[self->num_likelihood_nodes] = u; - self->num_likelihood_nodes++; -} - -/* 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) -{ - int j, k; - tsk_id_t *restrict L_nodes = self->likelihood_nodes; - - k = 0; - for (j = 0; j < self->num_likelihood_nodes; j++) { - L_nodes[k] = L_nodes[j]; - if (L_nodes[j] != node) { - k++; - } - } - assert(self->num_likelihood_nodes == k + 1); - self->num_likelihood_nodes = k; -} - -static inline void -ancestor_matcher_update_num_below(ancestor_matcher_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; - - for (; u != NULL_NODE; u = parent[u]) { - num_below[u] += delta; - } -} - -/* 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) -{ - const double *restrict L = self->likelihood; - double *restrict L_cache = self->likelihood_cache; - const tsk_id_t *restrict parent = self->parent; - double value; - 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); - } - value = L_cache[v]; - 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; - v = parent[v]; - } - return value; -} - -static void -ancestor_matcher_reset_likelihood_cache(ancestor_matcher_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) { - L_cache[u] = CACHE_UNSET; - u = parent[u]; - } -} - -/* Remove the specified edge 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). */ -static void -ancestor_matcher_remove_edge(ancestor_matcher_t *self, edge_t edge, - 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]); - } - if (update_likelihoods && L[c] == NULL_LIKELIHOOD) { - ancestor_matcher_add_likelihood( - self, c, ancestor_matcher_get_inherited_likelihood(self, p)); - } - remove_edge(edge, 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)) { - isolated[*num_isolated] = c; - (*num_isolated)++; - } - if (!ancestor_matcher_in_tree(self, p)) { - /* p was a root, and so has a likelihood */ - assert(L[p] >= 0); - isolated[*num_isolated] = p; - (*num_isolated)++; - } - } -} - -/* Insert the specified edge 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) -{ - 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[edge.child] == NOT_IN_TREE) { - ancestor_matcher_add_likelihood(self, edge.child, 0); - } - /* Otherwise the child was a root and keeps its likelihood */ - assert(L[edge.child] >= 0); - } - insert_edge(edge, 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]); - } -} - -/* 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; - -static void -tree_position_init( - tree_position_t *self, const matcher_indexes_t *matcher_indexes, int direction) -{ - 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; - tsk_id_t *restrict isolated = self->likelihood_nodes_tmp; - double *restrict L = self->likelihood; - int_fast32_t j, removal_start; - 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); - } - } - /* 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); - L[u] = NOT_IN_TREE; - } - } -} - -/* 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) -{ - const size_t N = self->num_nodes; - tsk_id_t u; - size_t j; - - 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)); - 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)); - - 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); - } else { - self->likelihood[u] = NOT_IN_TREE; - } - } -} - -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; - } -} - -/* 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) -{ - const double *restrict L = self->likelihood; - const tsk_id_t *restrict parent = self->parent; - mutation_list_node_t *mutation; - tsk_id_t u; - - for (mutation = self->matcher_indexes->sites.mutations[site]; mutation != NULL; - mutation = mutation->next) { - if (L[mutation->node] == NULL_LIKELIHOOD) { - u = mutation->node; - while (L[u] == NULL_LIKELIHOOD) { - u = parent[u]; - assert(u != NULL_NODE); - } - ancestor_matcher_add_likelihood(self, mutation->node, L[u]); - } - } -} - -/* Remove likelihoods that are equal to the likelihood inherited from the - * parent. */ -static void -ancestor_matcher_coalesce_likelihoods(ancestor_matcher_t *self) -{ - double L_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; - int j, num_cached_paths, num_likelihood_nodes; - - num_cached_paths = 0; - num_likelihood_nodes = 0; - for (j = 0; j < old_num_likelihood_nodes; j++) { - u = L_nodes[j]; - p = parent[u]; - 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 (L[u] >= 0) { - L_nodes[num_likelihood_nodes] = L_nodes[j]; - num_likelihood_nodes++; - } - } - self->num_likelihood_nodes = num_likelihood_nodes; - for (j = 0; j < num_cached_paths; j++) { - ancestor_matcher_reset_likelihood_cache(self, cached_paths[j]); - } -} - -static inline double -emission_probability(allele_t node_state, allele_t state, double mu, double num_alleles) -{ - if (state == TSLS_MISSING_DATA || node_state == state) { - return 1 - (num_alleles - 1) * mu; - } - return mu; -} - -static inline int -check_haplotype_state(const ancestor_matcher_t *self, tsk_id_t site, allele_t state) -{ - const tsk_size_t num_alleles = self->matcher_indexes->sites.num_alleles[site]; - - if (state < TSLS_MISSING_DATA || (state >= 0 && (tsk_size_t) state >= num_alleles)) { - return TSLS_ERR_BAD_HAPLOTYPE_ALLELE; - } - return 0; -} - -static inline double -ancestor_matcher_get_n(const ancestor_matcher_t *self) -{ - double n = 1.0; - if (!(self->flags & TSLS_DISABLE_WEIGHT_BY_N)) { - n = (double) self->num_nodes; - } - return n; -} - -/* 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) -{ - const double *restrict L = self->likelihood; - const tsk_id_t *restrict parent = self->parent; - 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; - tsk_id_t u, w; - int j; - - for (j = 0; j < self->num_likelihood_nodes; j++) { - u = L_nodes[j]; - num_inheriting[u] = num_below[u]; - } - for (j = 0; j < self->num_likelihood_nodes; j++) { - u = L_nodes[j]; - w = parent[u]; - if (w != NULL_NODE) { - while (L[w] == NULL_LIKELIHOOD) { - w = parent[w]; - } - num_inheriting[w] -= num_below[u]; - } - } -} - -/* Returns the sum of the likelihoods over all the nodes in the tree. */ -static double -ancestor_matcher_total_likelihood(const ancestor_matcher_t *self) -{ - const double *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; - tsk_id_t u; - int j; - - 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]; - } - return total; -} - -/* 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) -{ - int ret = 0; - tsk_id_t u; - int j; - int8_t *restrict list_R; - tsk_id_t *restrict list_node; - 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; - const int num_likelihood_nodes = self->num_likelihood_nodes; - bool match; - - /* Check to see if the previous site has the same recombination_required. If so, - * we can reuse the same list. */ - match = false; - if (site_id > 0) { - list = &T[site_id - 1]; - if (list->size == num_likelihood_nodes) { - list_node = list->node; - 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]]) { - match = false; - break; - } - } - } - } - - if (match) { - T[site_id].size = T[site_id - 1].size; - T[site_id].node = T[site_id - 1].node; - T[site_id].recombination_required = T[site_id - 1].recombination_required; - } else { - list_node = tsi_blkalloc_get(&self->traceback_allocator, - aligned_size((size_t) num_likelihood_nodes * sizeof(tsk_id_t))); - list_R = tsi_blkalloc_get(&self->traceback_allocator, - aligned_size((size_t) num_likelihood_nodes * sizeof(int8_t))); - if (list_node == NULL || list_R == NULL) { - ret = TSLS_ERR_NO_MEMORY; - goto out; - } - T[site_id].node = list_node; - T[site_id].recombination_required = list_R; - T[site_id].size = num_likelihood_nodes; - for (j = 0; j < num_likelihood_nodes; j++) { - u = nodes[j]; - list_node[j] = u; - list_R[j] = R[u]; - } - } - self->total_traceback_size += (size_t) num_likelihood_nodes; -out: - return ret; -} - -static int WARN_UNUSED -ancestor_matcher_process_site_viterbi(ancestor_matcher_t *self, const tsk_id_t site, - const allele_t state, compressed_matrix_t *output) -{ - 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; - 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]; - 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]; - 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_allelic_state(self, site); - max_L = -1; - 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_recomb = rho / n; - recombination_required[u] = false; - if (p_no_recomb > p_recomb) { - p_t = p_no_recomb; - } else { - p_t = p_recomb; - recombination_required[u] = true; - } - p_e = emission_probability( - 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); - (void) num_likelihood_nodes; - - if (max_L <= 0) { - if (mu <= 0 || mu >= 1) { - ret = TSLS_ERR_MATCH_IMPOSSIBLE_EXTREME_MUTATION_PROBA; - } else if (rho == 0) { - ret = TSLS_ERR_MATCH_IMPOSSIBLE_ZERO_RECOMB_PRECISION; - } else { - ret = TSLS_ERR_MATCH_IMPOSSIBLE; - } - goto out; - } - assert(max_L_node != NULL_NODE); - 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); - } - ancestor_matcher_prune(self); - ancestor_matcher_count_distinct(self); - ret = ancestor_matcher_store_traceback(self, site); - if (ret != 0) { - goto out; - } - if (output != NULL) { - ret = compressed_matrix_store_site(output, site, max_L, - self->num_likelihood_nodes, L_nodes, L, recombination_required); - if (ret != 0) { - goto out; - } - } - ancestor_matcher_coalesce_likelihoods(self); -out: - return ret; -} - -/* Recombination redistributes the total probability over the states at the - * previous site, including nodes that have since left the tree, so we - * 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, - const allele_t state, double *previous_total, compressed_matrix_t *output) -{ - 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; - 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]; - const double n = ancestor_matcher_get_n(self); - const double num_alleles = (double) self->matcher_indexes->sites.num_alleles[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); - total = *previous_total; - if (total < 0) { - total = ancestor_matcher_total_likelihood(self); - } - 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); - L[u] = (L[u] * (1 - rho) + rho / n * total) * p_e; - } - ancestor_matcher_unset_allelic_state(self, site); - normalisation_factor = ancestor_matcher_total_likelihood(self); - if (!(normalisation_factor > 0)) { - ret = TSLS_ERR_MATCH_IMPOSSIBLE; - goto out; - } - for (j = 0; j < self->num_likelihood_nodes; j++) { - 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); - 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); -out: - return ret; -} - -/* The backward transition from site l to site l - 1 depends on the tree at - * site l - 1 (nodes may have entered or left the tree), so we apply it - * lazily when we reach site l - 1, using these saved values from site l. - */ -typedef struct { - bool pending; - double total; - double rho; - double normalisation_factor; -} backward_transition_t; - -static int WARN_UNUSED -ancestor_matcher_process_site_backward(ancestor_matcher_t *self, const tsk_id_t site, - const allele_t state, const double *forward_normalisation_factor, - backward_transition_t *transition, compressed_matrix_t *output) -{ - 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; - 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); - const double num_alleles = (double) self->matcher_indexes->sites.num_alleles[site]; - double x; - tsk_id_t u; - int j; - - 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 / transition->normalisation_factor, self->likelihood_threshold); - } - ancestor_matcher_prune(self); - } - ancestor_matcher_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_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); - } - ancestor_matcher_unset_allelic_state(self, site); - transition->pending = true; - 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); -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, - const allele_t *haplotype, const double *forward_normalisation_factor, - compressed_matrix_t *output) -{ - int ret = 0; - 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; - tsk_id_t site, first_site; - double previous_total = -1; - - for (site = start; site < end; site++) { - ret = check_haplotype_state(self, site, haplotype[site]); - if (ret != 0) { - goto out; - } - } - 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); - 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); - } - if (self->flags & TSLS_EXTENDED_CHECKS) { - ancestor_matcher_check_state(self); - } - if (mode == MODE_VITERBI) { - ret = ancestor_matcher_process_site_viterbi( - self, site, haplotype[site], output); - } else if (mode == MODE_FORWARD) { - ret = ancestor_matcher_process_site_forward( - self, site, haplotype[site], &previous_total, output); - } else { - ret = ancestor_matcher_process_site_backward(self, site, haplotype[site], - forward_normalisation_factor, &transition, output); - } - if (ret != 0) { - goto out; - } - } -out: - 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) -{ - int j; - const int8_t *restrict R = self->traceback[site].recombination_required; - const tsk_id_t *restrict node = self->traceback[site].node; - const int size = self->traceback[site].size; - - for (j = 0; j < size; j++) { - recombination_required[node[j]] = R[j]; - } -} - -/* 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) -{ - int j; - const tsk_id_t *restrict node = self->traceback[site].node; - const int size = self->traceback[site].size; - - for (j = 0; j < size; j++) { - recombination_required[node[j]] = -1; - } -} - -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) -{ - tsk_id_t site; - edge_t edge; - tsk_id_t u, max_likelihood_node; - coordinate_t left, right, pos, start_pos, end_pos; - 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; - size_t path_length = 0; - - start_pos = start == 0 ? 0 : sites_position[start]; - end_pos = sites_position[end]; - - /* Prepare for the traceback and get the memory ready for recording - * the output edges. */ - path_right[path_length] = end_pos; - max_likelihood_node = self->max_likelihood_node[end - 1]; - assert(max_likelihood_node != NULL_NODE); - path_parent[path_length] = max_likelihood_node; - - /* Now go through the trees in reverse and run the traceback */ - 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); - } - 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); - } - } - } - path_left[path_length] = start_pos; - path_length++; - *path_length_out = path_length; -} - -static int -ancestor_matcher_reset_traceback(ancestor_matcher_t *self) -{ - int ret = 0; - - if (self->traceback_allocator.total_size > self->traceback_realloc_size) { - tsi_blkalloc_free(&self->traceback_allocator); - if (tsi_blkalloc_init(&self->traceback_allocator, self->traceback_block_size) - != 0) { - ret = TSLS_ERR_NO_MEMORY; - goto out; - } - } else { - tsi_blkalloc_reset(&self->traceback_allocator); - } - self->total_traceback_size = 0; - self->total_distinct_values = 0; -out: - return ret; -} - -/* Run the Viterbi algorithm and traceback. If output is not NULL, the - * normalised values and recombination_required flags stored in the - * 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, - 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 ret = 0; - - if (start < 0 || end > (tsk_id_t) self->num_sites || start >= end) { - ret = TSLS_ERR_BAD_SITE_RANGE; - goto out; - } - if (output != NULL && output->num_sites != self->num_sites) { - ret = TSLS_ERR_BAD_SITE_RANGE; - goto out; - } - ret = ancestor_matcher_reset_traceback(self); - if (ret != 0) { - goto out; - } - ret = ancestor_matcher_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, - 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)); - memset(self->allelic_state, 0xff, self->num_nodes * sizeof(*self->allelic_state)); - return ret; -} - -int -ancestor_matcher_find_path(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) -{ - return ancestor_matcher_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) -{ - int ret = 0; - - if (output->num_sites != self->num_sites) { - ret = TSLS_ERR_BAD_SITE_RANGE; - goto out; - } - if (self->num_sites > 0) { - ret = ancestor_matcher_run(self, MODE_FORWARD, 0, (tsk_id_t) self->num_sites, - haplotype, NULL, output); - } -out: - memset(self->allelic_state, 0xff, self->num_nodes * sizeof(*self->allelic_state)); - return ret; -} - -int -ancestor_matcher_backward(ancestor_matcher_t *self, const allele_t *haplotype, - const double *forward_normalisation_factor, compressed_matrix_t *output) -{ - int ret = 0; - size_t j; - - if (output->num_sites != self->num_sites) { - ret = TSLS_ERR_BAD_SITE_RANGE; - goto out; - } - for (j = 0; j < self->num_sites; j++) { - if (!(forward_normalisation_factor[j] > 0)) { - ret = TSLS_ERR_BAD_NORMALISATION_FACTOR; - goto out; - } - } - if (self->num_sites > 0) { - ret = ancestor_matcher_run(self, MODE_BACKWARD, 0, (tsk_id_t) self->num_sites, - haplotype, forward_normalisation_factor, output); - } -out: - memset(self->allelic_state, 0xff, self->num_nodes * sizeof(*self->allelic_state)); - return ret; -} - -double -ancestor_matcher_get_mean_traceback_size(ancestor_matcher_t *self) -{ - if (self->num_sites == 0) { - return 0; - } - return (double) self->total_traceback_size / ((double) self->num_sites); -} - -double -ancestor_matcher_get_mean_num_distinct_values(ancestor_matcher_t *self) -{ - if (self->num_sites == 0) { - return 0; - } - return (double) self->total_distinct_values / ((double) self->num_sites); -} - -size_t -ancestor_matcher_get_total_memory(ancestor_matcher_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..df5c043 100644 --- a/tsls/lib/tests/tests.c +++ b/tsls/lib/tests/tests.c @@ -13,129 +13,149 @@ 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_index_t index; + 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_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(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(index.mutation_node[mutation->id], mutation->node); + CU_ASSERT_EQUAL(index.mutation_derived_state[mutation->id], + acgt_code(mutation->derived_state[0])); + } + 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); - tsk_table_collection_free(&tables); + /* Per-site encoding, as in tskit */ + 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(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(index.mutation_derived_state[mutation->id], + mutation->derived_state[0] == 'A' ? 0 : 1); + } + tsls_index_free(&index); + + ret = tsls_index_init(&index, &ts, too_few, 2); + CU_ASSERT_EQUAL(ret, TSLS_ERR_UNKNOWN_ALLELE); + tsls_index_free(&index); + + ret = tsls_index_init(&index, &ts, acgt, TSLS_MAX_ALLELES + 1); + CU_ASSERT_EQUAL(ret, TSLS_ERR_TOO_MANY_ALLELES); + tsls_index_free(&index); + + 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_index_t index; + 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); + build_forest_ts(&ts, false); + ret = tsls_index_init(&index, &ts, acgt, 4); 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); + ret = tsls_hmm_alloc(&hmm, &index, 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 +164,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); - ret = ancestor_matcher_find_path( - &am, 2, 1, h, match, &path_length, left, right, parent); + /* 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 = 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 +198,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 +221,59 @@ 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); + tsls_index_free(&index); + 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_index_t index; + 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); + build_forest_ts(&ts, false); + ret = tsls_index_init(&index, &ts, acgt, 4); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = ancestor_matcher_alloc(&am, &mi, rho, mu, DBL_MIN, 0); + ret = tsls_hmm_alloc(&hmm, &index, 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); + tsls_index_free(&index); + tsk_treeseq_free(&ts); } static void @@ -296,7 +323,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 f823696..4ecde61 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,19 +92,85 @@ typedef struct { tsi_blkalloc_t allocator; } compressed_matrix_t; +/* 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 { + 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 */ + tsk_size_t *num_alleles; + allele_t *ancestral_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 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 { int flags; - const matcher_indexes_t *matcher_indexes; + 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; 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; @@ -145,14 +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; - 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; @@ -165,39 +196,37 @@ 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, - 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_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, + 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..14fa9aa 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,103 @@ 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, alleles=ALLELES) + assert tts.num_nodes == ts.num_nodes + assert tts.num_sites == ts.num_sites + m = ts.num_sites + 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) + 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,13 +583,12 @@ 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)) - am, _ = make_matcher(ts, 1e-2, 1e-2, True) + tsls.Hmm(ts, np.zeros(ts.num_sites), np.zeros(1)) + 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) @@ -529,10 +613,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 +668,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 +748,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..191104b 100644 --- a/tsls/tsls/matching.py +++ b/tsls/tsls/matching.py @@ -7,66 +7,31 @@ 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, 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. """ - 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, *, 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 - 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, - ) + self.alleles = alleles + super().__init__(ll_tables, alleles=alleles) class CompressedMatrix: @@ -181,11 +146,19 @@ 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. ``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, 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 @@ -200,16 +173,21 @@ 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, 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 pruning is not None: mode, parameter = pruning @@ -221,13 +199,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..2f9092f 100644 --- a/verification/verification.py +++ b/verification/verification.py @@ -146,13 +146,13 @@ 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, alleles=ALLELES) + return tsls.Hmm( + tree_sequence, np.full(m, rho), np.full(m, mu), weight_by_n=weight_by_n,