Validation plus experiments

.gitignore

@@ -224,3 +224,9 @@ __marimo__/
 # Benchmarking data and results
 benchmarking/data/
 benchmarking/results/
+
+# Verification output
+verification/tmp__NOBACKUP__/
+
+# Experiment results
+experiments/*/results/

README.md

@@ -25,7 +25,9 @@ submodule, so clone with `--recurse-submodules` or run
 module `_tsls`). It is based on the tsinfer engine:
 
 - It copies from every node in the tree at a given position (any node
-  with a parent or a child); isolated nodes cannot be copied.
+  with a parent or a child); isolated nodes cannot be copied. Copying a
+  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
@@ -39,7 +41,8 @@ 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
 isolated node 0. The `benchmarking/` directory compares the engines on
-the sc2ts ARG.
+the sc2ts ARG, and `verification/` runs statistical validations of tsls on
+simulated ARGs (plots, not pass/fail tests).
 
 The git history records exactly what was copied from where. Each
 directory was first added verbatim, then changed in separate commits to

experiments/thresholding/README.md

@@ -0,0 +1,144 @@
+# Thresholding experiments
+
+`thresholding.py` measures how well different ways of limiting the number of
+stored likelihood values work in the tsls engine. It also measures what each
+one costs on the tasks from `verification/`.
+
+```
+uv sync
+uv run python thresholding.py --quick -t 4    # about 2 minutes
+uv run python thresholding.py -t 4            # full grid; about an hour
+uv run python thresholding.py --analyse-only  # redo plots and summary
+```
+
+Results, plots and `summary.md` are written to `results/`, which is not
+tracked by git. The sc2ts scenario needs the sub-ARGs and queries from
+`benchmarking/data`; if they are missing, the scenario is skipped.
+
+## Methods
+
+Every method is applied after the values are normalised at each site:
+by the maximum for Viterbi, and by the sum for forward/backward. A "pruned"
+node is given the value `likelihood_threshold`, which by default is the
+smallest positive double. It then merges into its parent, and can only be
+reached again by recombination.
+
+- **`floor`** is the existing `likelihood_threshold`: any value below `t` is
+  set to `t`.
+- **`relative`** prunes nodes with a value below `f` times the maximum. Use
+  `pruning=("relative", f)` (experimental option).
+- **`top_n`** prunes all but the `n` nodes with the largest values. Nodes tied
+  with the n-th value are kept. Use `pruning=("top_n", n)` (experimental option).
+- **`lossless_floor`** is `floor` at `t* = (ρ/n) / (1 − ρ + ρ/n)`. In Viterbi,
+  any node below `t*` loses to recombination at the next site, so this floor
+  changes no Viterbi likelihoods. With `weight_by_n=False`, as sc2ts uses,
+  `t* = ρ`.
+
+`AncestorMatcher(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`).
+
+## Scenarios
+
+- **threading**: a held-out sample from an msprime simulation, matched as in
+  `verification.py` (`HeldOutThreading`). The HMM uses sc2ts parameters
+  (k = 3, `weight_by_n=False`). For Viterbi we measure the loss in path
+  log-likelihood relative to exact, how often the path is unchanged, and the
+  accuracy of copying the true parent. For forward/backward we measure the
+  posterior total-variation distance from exact and the posterior on the
+  true parent.
+- **calibration**: queries simulated from the HMM itself, as in
+  `verification.py` (`ModelCalibration`). For forward/backward we measure
+  the reliability error, and 90% credible-set coverage minus the expected
+  coverage.
+- **sc2ts**: queries from `benchmarking/` on sc2ts sub-ARGs, using the tsls
+  benchmark engine settings. We measure the loss in Viterbi path
+  log-likelihood, how often the path is unchanged, time, and stored values.
+  At 10k samples we also measure the forward/backward posterior
+  total-variation distance, at 100 evenly spaced sites.
+
+## Findings (quick run)
+
+These numbers come from `--quick`, which uses:
+- 3 replicates of threading, with 50 samples only;
+- 3 ARGs per calibration setting;
+- 3 sc2ts queries on the 10k sub-ARG (an exact copy, and recombinants with 2
+  and 5 extra mutations);
+- a reduced grid: t and f in {1e-20, 1e-8, 1e-4, 1e-2}, and n in
+  {10, 100, 1000}.
+
+Treat them as indicative. The full grid, with 50–1000 samples and sc2ts
+sub-ARGs up to 100k, was not run to completion.
+
+### 1. For Viterbi, compression limits storage more than thresholding does
+
+On the 10k sc2ts sub-ARG, exact Viterbi stores 312–1,248 nodes per site,
+but only **4–6 distinct values** per site. With 50 simulated samples, it
+stores 9.4 nodes for 4.4 distinct values.
+
+The tsinfer-style compression only removes a node's value if it equals its
+parent's. Siblings with the same value are kept separately, so most stored
+nodes repeat a few values. Thresholding changes values, not structure. The
+best any method achieved without loss was about a 2× smaller traceback:
+
+| sc2ts 10k, Viterbi | traceback vs exact | time vs exact | log-likelihood loss | identical paths |
+|---|---|---|---|---|
+| `lossless_floor` | 1.00 | 0.96 | 0 | 3/3 |
+| `top_n` = 10 | 0.52 | 0.61 | 0 | 3/3 |
+| `relative` = 1e-4 | 1.00 | 1.04 | 0 | 3/3 |
+| `floor` = 1e-4 | 0.93 | 0.90 | 0.06 | 1/3 |
+| `floor` = 1e-2 | 0.44 | 0.36 | 1.5 | 1/3 |
+| `relative` = 1e-2 | 0.44 | 0.34 | 5.9 | 1/3 |
+
+Merging equal-valued siblings, which is closer to tskit's parsimony step but
+cheaper, looks like the larger opportunity.
+
+### 2. For Viterbi, a floor above `t*` is harmful, and pruning should be used instead
+
+A floor raises weak nodes up to `t`. Once `t` is above the recombination
+value, staying on a raised node beats switching, so stale copying targets
+survive. Pruning sets weak nodes close to 0 instead, so they lose to
+recombination as they should.
+
+| threading (50 samples), Viterbi | traceback vs exact | log-likelihood loss | accuracy drop |
+|---|---|---|---|
+| `floor` = 1e-4 | 0.92 | 200 | 0.34 |
+| `relative` = 1e-4 | 0.92 | 0 | 0 |
+| `top_n` = 10 | 0.98 | 0 | 0 |
+| `lossless_floor` (2.4e-6) | 0.99 | 0 | 0 |
+
+The lossless floor explains why every threshold up to 1e-6 was free in
+earlier runs: with sc2ts settings, `t* = ρ ≈ 2.4e-6`.
+
+### 3. For forward/backward, every method distorts the posteriors, in different directions
+
+- **`floor`** works by adding probability to every node, so posteriors no
+  longer sum to 1 (they are off by about 0.99 at 50 samples). It is still
+  the least biased of the three.
+  - On sc2ts 10k, `floor` = 1e-4 cut the stored forward values by about
+    200× and the forward-backward time to under 1%. The posterior
+    total-variation distance from exact was 0.02.
+  - With 50 samples the same setting gave a total-variation distance of
+    0.07, and calibration was as good as exact.
+  - `floor` = 1e-2 made the posteriors under-confident: 90% coverage was
+    9 points above expected.
+- **`relative`** and **`top_n`** make the posteriors over-confident:
+  - `relative` = 1e-4 gave a total-variation distance of 0.44 with 50
+    samples;
+  - `top_n` = 10 gave 0.13 with 50 samples and 0.68 on sc2ts;
+  - 90% coverage was 7–28 points below expected.
+
+  With `top_n` = 100 on sc2ts, more values were stored than exact (1.23×),
+  because pruned values stop compressing as before.
+
+## Recommendations
+
+- **Viterbi:** use the lossless floor `t*` as the default threshold; it
+  costs nothing. If more speed is needed, use `relative` or `top_n` rather
+  than a floor above `t*`.
+- **Forward/backward:** treat any thresholding as approximate. A floor
+  around 1e-4 gives large speedups on sc2ts with small posterior error.
+  Check the posterior total-variation distance for the intended use.
+- **Next step:** improve compression by merging siblings with equal values,
+  rather than adding more thresholding. Use `count_distinct=True` to measure
+  the potential gain.

experiments/thresholding/pyproject.toml

@@ -0,0 +1,27 @@
+[project]
+name = "tsls-thresholding-experiments"
+version = "0.0.0"
+description = "Experiments on likelihood thresholding and pruning in tsls."
+requires-python = ">=3.11"
+dependencies = [
+    "matplotlib",
+    "msprime",
+    "numpy>=2",
+    "pandas",
+    "tabulate",
+    "tqdm",
+    "tskit>=1.0.0",
+    "tsls",
+]
+
+[tool.uv]
+package = false
+
+[tool.uv.sources]
+tsls = { path = "../../tsls", editable = true }
+
+[tool.ruff]
+line-length = 89
+
+[tool.ruff.lint.isort]
+known-first-party = ["dense", "tsls", "verification"]