Some work on improving HMM engine performance

benchmarking/README.md

@@ -363,10 +363,84 @@ for the whole group divided by the number of haplotypes.
   not yet in the tree sequence.
 - tsls is 3–10% slower than tsinfer-hmm in mean time per query, with the
   same matcher memory (64 MiB for ancestors, about 90 MiB for samples)
-  and traceback sizes. Building tsls's `MatcherIndexes` takes several
+  and traceback sizes. (These figures predate the engine changes
+  described below, and `run.py` timings are biased by run order; see
+  [tsls against tsinfer-hmm](#tsls-against-tsinfer-hmm-where-the-time-goes).)
+  Building tsls's `MatcherIndexes` takes several
   times longer (3.6 s against 0.4 s for the full ancestors tree
   sequence), but this is done once per group.
 - Per-query times are close to tsinfer's per-haplotype CPU time, so the
   HMM accounts for most of tsinfer's matching cost.
 - The cost is in the last few ancestor groups and the samples. Matching
   all 6,404 samples takes about 4.5 CPU hours with either engine.
+
+### tsls against tsinfer-hmm: where the time goes
+
+tsls now stores a smaller Viterbi traceback than tsinfer-hmm (13 against
+21 entries per site for samples; see `experiments/compression/README.md`),
+but it isn't faster, because storing the traceback is under 1% of the
+time in both engines. Both engines update the same likelihood lists at
+each site, since tsinfer-hmm coalesces them right after storing. The time
+goes into walking up the tree, to find each listed node's allele and the
+value it inherits, and into moving from tree to tree. Building the first
+tree of a query means processing most of the edges before its window,
+which matters for late ancestors.
+
+Instruction counts (callgrind, collection limited to
+`ancestor_matcher_find_path`, 20 group 1485 ancestors and 1 sample) showed
+where tsls did more work:
+
+| | tsinfer-hmm | tsls before | tsls after |
+|---|---|---|---|
+| 20 ancestors, group 1485 | 6.65 G | 9.89 G | 7.44 G |
+| 1 sample | 9.92 G | 12.50 G | 10.96 G |
+
+Two changes to `tsls/lib/matcher.c` account for the difference between
+"before" and "after":
+
+1. **Allele lookup.** tsinfer-hmm's trees have the single root 0, whose
+   allele it sets to the ancestral state, so the walk up the tree only
+   checks for a set allele. tsls allows several roots, so
+   `get_allelic_state` also checked for the root at each step. tsls now
+   sets every root's allele to the ancestral state before the site's
+   update (the roots all have likelihoods, so they are found in the
+   likelihood list), and the three likelihood loops use
+   `get_marked_allelic_state`, which walks without the check. The walk
+   deliberately has no `assert`: the Python module is compiled with
+   asserts on, and an `assert` there cost as much as the check it
+   replaced (12.71 G on the sample, against 12.50 G before). The traceback
+   keeps the general `get_allelic_state`, since its tree has no
+   likelihoods.
+2. **Building the first tree.** For Viterbi, `init_sweep` now builds the
+   first tree with plain edge insertions and removals
+   (`ancestor_matcher_next_tree_topology`), as tsinfer-hmm does, instead
+   of passing every edge through the general wrappers that handle
+   likelihoods, subtree counts and isolated nodes. It also no longer
+   initialises `num_below`, which only forward and backward use. Forward
+   and backward still build it the general way, because they track
+   `num_below`.
+
+Wall times, measured in pairs: each query is run three times with each
+engine, alternating, and the best time is kept. Ratios are total tsls time
+over total tsinfer-hmm time (100 group 1485 ancestors, 6 samples):
+
+| | tsls before ÷ tsinfer-hmm | tsls after ÷ tsinfer-hmm |
+|---|---|---|
+| ancestors, group 1485 | 1.35 | 1.10 |
+| samples | 1.06 | 1.06 |
+
+- Change 2 is a clear win on late ancestors.
+- Change 1 removes about 1.5 G instructions per sample but no wall time.
+  The walk up the tree is limited by memory access (each step loads
+  `allelic_state` and `parent` for a node that is usually not in cache),
+  so the extra compare was nearly free.
+- What is left is mostly the same kind of overhead in the coalescing and
+  cache code: `ancestor_matcher_get_inherited_likelihood` and
+  `ancestor_matcher_reset_likelihood_cache` have asserts and root checks
+  inside their walks, which tsinfer-hmm's equivalents don't. By the
+  result for change 1, these would also save instructions more than time.
+- Timing caveat: `run.py` runs all queries for one engine, then all for
+  the next, so drift in machine speed (large on this VM: the same query
+  can vary by 20% between passes) shows up as a difference between
+  engines. Ratios from `run.py` ranged from 1.09 to 1.25 on the same
+  builds. Use paired measurements, as above, to compare engines.

experiments/compression/README.md

@@ -11,6 +11,10 @@ uv run python capture.py --sizes 10000,30000,100000   # about 35 minutes
 uv run python report.py                               # results/report.md
 ```
 
+`capture.py --cases` captures a subset of the cases (`msprime`, `sc2ts`,
+`tsinfer`), replacing their rows in an existing manifest and keeping the
+others. `report.py --cases` analyses a subset, for quick iteration.
+
 `capture.py` uses `AncestorMatcher.viterbi_matrix` and `forward_matrix`. It
 saves each result as a `CompressedMatrix` with `CompressedMatrix.dump`, next
 to the tree sequence it refers to. Reload a matrix with
@@ -30,6 +34,21 @@ The cases captured are:
   recombinants with one or two breakpoints.
 - **msprime ARGs**: held-out samples matched against 50 or 1000 other
   samples, with the same HMM settings.
+- **tsinfer workloads**: the late ancestor matching groups 1483, 1484 and
+  1485, and sample matching, from the tsinfer 1.0 inference of 1000 Genomes
+  chr20p (580,909 sites). These need the workloads made by
+  `benchmarking/tsinfer_data.py` (see `benchmarking/README.md`). They use
+  tsinfer's settings (ρ = 0.01, μ = 1e-20, `weight_by_n=True`, two
+  alleles). For each group there are 4 ancestors at spread quantiles of
+  window length, and there are 4 random sample haplotypes. Each path is
+  checked against the one tsinfer found.
+  - Viterbi runs over each query's window: the whole genome for samples.
+  - `forward_matrix` has no start and end, so forward runs on the tree
+    sequence restricted to the window's sites (`window_ts`), which keeps the
+    node IDs and n. For samples the forward window is the middle 20,000
+    sites (`--sample-forward-sites`). `capture.py --quick` checks the
+    windowing by comparing windowed Viterbi with Viterbi over
+    `[start, end)`.
 
 `analysis.py` computes the measures:
 - From the stored lists, at every site: the list size, distinct values,
@@ -42,6 +61,16 @@ The cases captured are:
   - change points after rounding the values;
   - an approximate minimum if only samples had to be represented, as in
     tskit (10 sites).
+- Encoding measures, comparing each site with the previous one:
+  - churn: nodes entering or leaving the stored list;
+  - changed values of nodes in both lists (a delta encoding);
+  - affine residuals: values of nodes in both lists that the site's global
+    map does not predict. For Viterbi the map is
+    `f(v) = max(v·stay, ρ/n)·e/c`, with `stay = 1 − ρ + ρ/n`, `e` the
+    matching emission and `c` the normalisation factor; this is exact. For
+    forward it is `f(v) = a·v + b`, with `a = (1 − ρ)·e/c` and `b` fitted.
+  - the same on the minimal (change point) lists at the decoded sites,
+    whose previous sites are decoded too.
 
 ## The representation
 
@@ -87,6 +116,11 @@ the site that end up equal to their parent.
 Representing only the samples, as tskit does, would save at most 5–30% on
 sc2ts. So the all-nodes model is not what makes the lists large.
 
+This does not hold on tsinfer workloads, where the trees change much more
+often. There, a quarter to a third of the Viterbi entries were redundant
+until the engine was changed to coalesce before storing. See
+[Affine and delta encodings on tsinfer workloads](#affine-and-delta-encodings-on-tsinfer-workloads).
+
 ### 2. Why Viterbi stores so many nodes with so few values on sc2ts
 
 The Viterbi values are almost entirely four exact classes:
@@ -155,3 +189,127 @@ where `floor = 1e-4` cut forward storage by about 200×.
   approximation.
 - The captured matrices (`results/*.npz`) can be used to size and check
   any new encoding offline before it is implemented in C.
+
+## Affine and delta encodings on tsinfer workloads
+
+Means over 4 queries per workload, for late ancestor matching (groups
+1483–1485) and sample matching on 1000 Genomes chr20p. Viterbi covers each
+query's whole window; forward covers ancestor windows and 20,000 sites for
+samples.
+
+### Before coalescing the traceback
+
+| | ancestors g1483 | g1484 | g1485 | samples |
+|---|---|---|---|---|
+| Viterbi stored / site | 12.2 | 11.5 | 20.1 | 26.9 |
+| Viterbi change points / site | 9.1 | 8.3 | 13.2 | 19.8 |
+| Viterbi churn (entering + leaving) / site | 6.0 | 6.3 | 13.8 | 16.1 |
+| Viterbi values changed / site | 0.73 | 0.76 | 1.13 | 1.25 |
+| Viterbi affine residuals / site | 0.22 | 0.19 | 0.32 | 0.28 |
+| Traceback: engine ÷ delta encoding | 1.6× | 1.6× | 1.1× | 1.7× |
+| Forward stored / site | 73 | 80 | 139 | 185 |
+| Forward values changed / site | 70 | 77 | 132 | 176 |
+| Forward affine residuals / site | 1.6 | 1.6 | 2.8 | 3.1 |
+| Forward: stored ÷ affine + delta | 7.8× | 8.5× | 8.0× | 8.5× |
+
+- Viterbi lists are small (12–27 per site, against hundreds to thousands
+  on sc2ts), and their values hardly change: about one per site, and only
+  0.2–0.3 per site that the global map `max(v·stay, ρ/n)·e/c` doesn't
+  predict. `weight_by_n=True` doesn't change this, because the uniform stay
+  factor cancels in normalisation. Exact and tolerant comparisons agree, so
+  rounding doesn't break a delta encoding.
+- But the lists churn: half or more of the entries enter or leave at each
+  site, because the tree changes about every second site. So delta
+  encoding the traceback saved only 1.1–1.7×, against 50–430× on sc2ts.
+- Most of the churn was the engine's bookkeeping, not structure in the
+  likelihoods. When edges change, nodes get explicit copies of their
+  inherited values, and the engine stored the traceback before removing
+  values equal to their parent's. A quarter to a third of the stored
+  entries were redundant (7 of 27 for samples), against almost none on
+  sc2ts. On the
+  minimal lists the churn was only 1.3–2.2 per site.
+- Forward values all change at every site, so delta encoding doesn't help
+  (0.95×), but the affine map `a·v + b` predicts all but 1.6–3.1 values per
+  site. Forward has the same redundant entries (3–9 per site), because it
+  also stores before coalescing.
+
+### Coalescing before storing the traceback
+
+The fix is to coalesce before storing. The catch is the traceback: it
+stores only `(node, recombination_required)`, and during traceback a node
+not in the list takes the flag of its nearest ancestor in the list. A node
+with the same value as its parent but a different flag can't be removed
+without changing the path. So the coalescing pass, when run before storing
+in Viterbi, removes a node only if both its value and its flag equal the
+ones it inherits. Nodes kept because of their flag are removed at a later
+site, once their flag is recomputed.
+
+What was tried, in order:
+
+1. **Flag-aware coalescing before storing, plus the old pass after
+   storing.** The stored lists became exactly the change points, and the
+   paths were unchanged. But it ran two coalescing passes per site, and tsls
+   went from about 1.1× to 1.36–1.48× tsinfer-hmm's time.
+2. **Flag-aware coalescing before storing only (one pass per site).** The
+   stored lists are the same as in 1. The nearest-ancestor walk for the
+   flag cost 2–4% more instructions than the old engine (table below). Its
+   time relative to tsinfer-hmm, measured in separate runs, looked like
+   1.13–1.22× (against 1.09–1.10× before). This turned out to be mostly
+   noise in comparing ratios across runs.
+3. **As 2, with the inherited flag cached** (landed). The flag is cached
+   along tree paths with the inherited likelihood
+   (`ancestor_matcher_get_inherited_likelihood_and_flag`, which uses the new
+   `recombination_required_cache` array), so the walk is shared. It costs
+   no more than the old engine.
+
+Instruction counts, from callgrind with collection limited to
+`ancestor_matcher_find_path`, so they are deterministic:
+
+| | old engine | 2: flag lookup uncached | 3: flag lookup cached |
+|---|---|---|---|
+| 20 ancestors, group 1485 | 9.92 G | 10.13 G (+2.2%) | 9.89 G (−0.2%) |
+| 1 sample | 12.54 G | 13.01 G (+3.8%) | 12.50 G (−0.3%) |
+
+Head-to-head wall times of the old engine and 3 agree within noise
+(4 samples: 11.2 vs 11.3 s; 100 group 1485 ancestors: 14.1 vs 14.4 s, best
+of several runs, in alternating order). The 1.36–1.48× for version 1 is
+also a ratio to tsinfer-hmm from a separate run. It is well outside the
+noise, but not as precise as the counts above.
+
+The matrices below were captured with version 1; versions 2 and 3 store
+the same lists.
+
+| | ancestors g1483 | g1484 | g1485 | samples |
+|---|---|---|---|---|
+| Viterbi stored / site | 9.2 | 8.3 | 13.1 | 18.8 |
+| Viterbi churn / site | 1.3 | 1.3 | 1.8 | 2.2 |
+| Viterbi values changed / site | 0.23 | 0.21 | 0.32 | 0.26 |
+| Engine traceback entries (change) | −35% | −35% | −47% | −31% |
+| Traceback: engine ÷ delta encoding | 5.1× | 5.0× | 5.2× | 8.4× |
+| Traceback: old engine ÷ delta encoding | 7.8× | 7.7× | 9.7× | 12× |
+
+- The stored lists are now exactly the change points, and churn falls about
+  5–8×. Paths are unchanged: tsls still finds tsinfer's path for every
+  query in all six benchmark workloads (2,014 queries).
+- The engine's traceback shrinks by a third to a half, and a delta encoding
+  on top of it saves another 5–8×. In the benchmarks, mean traceback size
+  for samples fell from 21.6 to 13.4, and the matcher's memory from 91 to
+  64–71 MiB.
+- Not changed: forward and backward still coalesce after storing (values
+  only; they have no flags), so their stored lists keep the redundant
+  entries described above.
+- The sc2ts and msprime results earlier in this README were captured with
+  the old engine. They had almost no redundant entries, so they shouldn't
+  change much.
+
+### Implications for tsinfer
+
+- Coalesce before storing (done for Viterbi) before trying a delta encoding.
+  Delta encoding then cuts traceback memory by 5–8×, though tsinfer's
+  traceback is already small (13 entries per site for samples).
+- In an affine Viterbi engine, the per-site work would be dominated by
+  churn from edge insertions and removals (1–2 per site), not by value
+  updates (0.1–0.3 per site).
+- For forward and backward, the affine update is the clear win: about 8×
+  fewer explicit updates on tsinfer workloads, and 50–100× on sc2ts.
+  Forward should also coalesce before storing.