Perf batch

benchmarking/README.md

@@ -745,3 +745,136 @@ time in proportion to all 3.4M nodes, so most trees aren't numbered.
 Numbering 16 times more eagerly made more sites sparse (6,490 and 9,780
 for queries 11 and 22) but was slower (35 and 40 s), so numbering whole
 trees is the limit at this scale.
+
+## Walk caches: sample matching under a recombination map (October 2026)
+
+Under the HapMap II chr20 map, Viterbi on a chr20p sample took about
+4.1 s against 2.2 s with tsinfer's constant r, because values from
+different sites' switch odds rarely coalesce, so the likelihood list is
+about 8 times longer (`experiments/coarsening`). This round assessed
+algorithmic strategies, with sample matching under the map as the main
+target, and stopped when every remaining idea's estimated gain was under
+10%. Each step's commit records its measurements.
+
+### Measuring
+
+- `profile_samples.py` runs tsls on chr20p samples under a map or
+  constant r, printing times and a hash of each path.
+- `ab.py` compares builds of the extension (each a directory holding
+  `_tsls*.so`, put first on `PYTHONPATH`): each round runs every build in
+  its own process, pinned to one core, in turn, keeping the best time per
+  query over the rounds, and checks that the paths are identical.
+- Where the time goes was found with `sudo perf record` and
+  `perf annotate -l` on `tsls_hmm_find_path_matrix`, into which the site
+  update, coalescing and walks are all inlined, so source lines rather
+  than functions give the split. Step counters (walk lengths, edges per
+  site, how often caches could survive) were added temporarily and not
+  committed.
+
+### Where the time went
+
+At the start, the time was in walks up the tree, not in the values:
+
+| | map | constant r |
+|---|---|---|
+| allele walks up to a mutation or root, and clearing their marks | ~30% | ~27% |
+| coalescing walks to the inherited value, and cache writes | ~24% | ~22% |
+| walks resetting the caches | ~14% | ~15% |
+| arithmetic on the values | <1% | <1% |
+
+(Shares of `tsls_hmm_find_path_matrix`, which was about 80% of the run
+under the map and 63% under constant r; `tsls_hmm_update_tree` was most
+of the rest.) There were about 7.2 edges out and 7.2 in per site.
+
+### Changes
+
+| commit | change | map | constant r |
+|---|---|---|---|
+| `009f587` | clear the caches with a stamp, not reset walks | 0.93 | 0.98 |
+| `a1d957f` | get alleles, and coalesce, through each listed node's nearest listed ancestor | 0.78 | 0.86 |
+| `f00bcc0` | mark listed nodes, so walks read only node records | 0.86 | 0.92 |
+| `39fd01e` | keep the tree as parent links and child counts | 0.96 | 0.93 |
+| `7914bc2` | keep the walk caches across sites | 0.72 | 0.70 |
+
+Ratios are paired wall times over the previous build (6 samples, best of
+3). The two main ideas:
+
+- **The list's own tree.** Every node with a mutation at a site is in
+  the likelihood list, so a listed node's allele is that of its nearest
+  listed ancestor. One walk per listed node, to its nearest listed
+  ancestor, serves both the alleles and coalescing, and these walks are
+  much shorter than walks to a mutation or root.
+- **Persistent walk caches.** Each walk caches its answer, a node ID, on
+  the unlisted nodes it passes. The answer stays right until a node on a
+  cached path gets a likelihood or a new parent, which happened at 24% of
+  sites under the map and 14% under constant r. A node that loses its
+  likelihood becomes a cached node pointing to its old listed ancestor.
+  `tsls_hmm_check_state` checks that every cached walk leads to the
+  nearest listed ancestor.
+
+Overall, against the build before these changes, with identical paths:
+
+| workload | time |
+|---|---|
+| samples, map, mu = 1e-20 | 0.45-0.47 |
+| samples, map, mu = 1e-8 | 0.46 |
+| samples, constant r | 0.50-0.52 |
+
+`paired.py` (tsls over tsinfer-hmm; the path classes are unchanged):
+
+| workload | before | after |
+|---|---|---|
+| ancestors, group 84 | 0.773 | 0.674 |
+| ancestors, group 1485 | 0.714 | 0.528 |
+| sc2ts 100k, k = 1 | 0.020 | 0.018 |
+| sc2ts 100k, exact | 0.013 | 0.012 |
+
+A sample now takes about 2.0-2.3 s under the map and 0.8-0.9 s under
+constant r.
+
+### Dead ends and remaining ideas
+
+- **Reusing listed nodes' own ancestors between sites** (stamped, and
+  recomputed when stale): 1.00-1.02. Checking a stored ancestor costs a
+  record load, as the one-step walk it replaces does.
+- **Batching queries** in one sweep: only the tree position stepping and
+  the structural part of the edge updates are shared (about 7% of the run
+  after the changes above, and at most about 15% before its own
+  overheads), since each query's list, walks and caches are its own.
+  Not pursued.
+- **Remaining profile** (map, shares of the run): walks to listed
+  ancestors ~22%, tree updates ~13%, tree position stepping ~7%, the
+  site update's arithmetic ~7%, the traceback's walk up from the path
+  node to find the matched allele ~6%, storing the traceback ~5%, the
+  allele walk over the list ~4%, inserting mutation nodes ~4%. The
+  candidates each have a ceiling under 10%: precomputing tree transitions
+  in the index (~5-7%), bounding the traceback's allele walk by the
+  mutation's time (exact, since parents are older than children; ~6%),
+  and fusing traceback storage into the site update (~5%).
+- Approximations were not needed; earlier thresholding and coarsening
+  work found them costly in accuracy for little time.
+
+### Against tsinfer-hmm after the walk caches
+
+Two full rounds of `paired.py` (`9b0112b`, pinned to one core with
+`taskset`, best of 3 per query), tsls / tsinfer-hmm. `samples_map` and
+`samples_map_mu1e-8` are the 6 samples under the HapMap II chr20 map.
+
+| workload | queries | round 1 | round 2 | memory (MiB), tsls / tsinfer-hmm | paths same/tie/better/worse |
+|---|---|---|---|---|---|
+| ancestors, group 84 | 50 | 0.64 | 0.67 | 4 / 64 | 50/0/0/0 |
+| ancestors, group 1000 | 50 | 0.51 | 0.52 | 4 / 64 | 50/0/0/0 |
+| ancestors, group 1483 | 100 | 0.46 | 0.49 | 4 / 64 | 100/0/0/0 |
+| ancestors, group 1484 | 50 | 0.39 | 0.47 | 8 / 64 | 50/0/0/0 |
+| ancestors, group 1485 | 100 | 0.51 | 0.50 | 4 / 64 | 100/0/0/0 |
+| samples, constant r | 6 | 0.30 | 0.30 | 12 / 128 | 6/0/0/0 |
+| samples, map, mu = 1e-20 | 6 | 0.16 | 0.17 | 36 / 384 | 6/0/0/0 |
+| samples, map, mu = 1e-8 | 6 | 0.17 | 0.16 | 40 / 448 | 6/0/0/0 |
+| sc2ts 100k, k = 1 | 10 | 0.018 | 0.018 | 4 / 640 | 4/2/3/1 |
+| sc2ts 100k, exact | 10 | 0.013 | 0.013 | 4 / 960 | 8/2/0/0 |
+| sc2ts 300k, exact | 4 | 0.044 | 0.043 | 4 / 1,792 | 3/1/0/0 |
+
+A chr20p sample takes about 0.8 s with tsls against 2.7 s with
+tsinfer-hmm under constant r, and 1.7-1.9 s against 10-12 s under the
+map. The one worse sc2ts k = 1 path is the tie in floored values
+described in `c81ac3b`.

benchmarking/ab.py

@@ -0,0 +1,68 @@
+"""
+Compare the wall time of two or more builds of the tsls extension on chr20p
+sample matching, with paired runs.
+
+Each build is a directory holding a _tsls extension module, put first on
+PYTHONPATH ("." means the in-tree build). For each round, every build runs
+profile_samples.py in its own process, pinned to one core, in turn; the
+best time for each (build, query) over the rounds is kept. Reports the
+total time of each build over the first, and checks that every build finds
+the same paths.
+
+    uv run python ab.py base=/path/to/build new=. --map data/genetic_maps/... -r 3
+"""
+
+import collections
+import json
+import os
+import subprocess
+import sys
+
+import click
+
+
+def run_build(path, args, core):
+    env = dict(os.environ)
+    if path != ".":
+        env["PYTHONPATH"] = path
+    command = ["taskset", "-c", str(core), sys.executable, "profile_samples.py", *args]
+    out = subprocess.run(command, env=env, capture_output=True, text=True, check=True)
+    return [json.loads(line) for line in out.stdout.splitlines() if line.startswith("{")]
+
+
+@click.command()
+@click.argument("builds", nargs=-1, required=True)
+@click.option("--map", "genetic_map", default=None)
+@click.option("--mu", type=float, default=1e-20)
+@click.option("--queries", "query_ids", default="0,40,80,120,160,199")
+@click.option("-r", "--rounds", type=int, default=3)
+@click.option("--core", type=int, default=2)
+def main(builds, genetic_map, mu, query_ids, rounds, core):
+    named = [b.split("=", 1) for b in builds]
+    args = ["--mu", str(mu), "--queries", query_ids]
+    if genetic_map is not None:
+        args += ["--map", genetic_map]
+    best = collections.defaultdict(dict)
+    paths = collections.defaultdict(dict)
+    for round_ in range(rounds):
+        for name, path in named:
+            for row in run_build(path, args, core):
+                q = row["query"]
+                best[name][q] = min(best[name].get(q, float("inf")), row["time"])
+                paths[name][q] = row["path"]
+        totals = {name: sum(best[name].values()) for name, _ in named}
+        base = named[0][0]
+        click.echo(
+            f"round {round_ + 1}: "
+            + ", ".join(
+                f"{name} {totals[name]:.2f}s ({totals[name] / totals[base]:.3f})"
+                for name, _ in named
+            )
+        )
+    for name, _ in named[1:]:
+        same = paths[name] == paths[named[0][0]]
+        click.echo(f"{name}: paths {'identical' if same else 'DIFFER'}")
+
+
+if __name__ == "__main__":
+    main()

benchmarking/paired.py

@@ -11,7 +11,9 @@
 There are two kinds of workload:
 
 - tsinfer: the chr20p ancestor groups and samples made by tsinfer_data.py,
-  with the model in their model.json.
+  with the model in their model.json. The "samples_map" workloads replace
+  its constant rho with per-site probabilities from the HapMap II chr20 map
+  (genetic_map.py), and "_mu1e-8" also replaces its mu.
 - sc2ts: queries from data/queries_scaling matched against the sc2ts
   sub-ARGs, with the sc2ts model (num_mismatches = 3, weight_by_n=False).
   sc2ts matches every sample first with a likelihood threshold of k = 0
@@ -32,6 +34,7 @@
 
 import click
 import numpy as np
+import tsls
 
 import callgrind
 import engines
@@ -53,11 +56,18 @@ class Workload:
     # None for tsinfer workloads, which use their model.json
     mismatch_threshold: float = None
     sc2ts: bool = False
+    # A genetic map (relative to the data directory) and mu to use in place
+    # of model.json's rho and mu
+    genetic_map: str = None
+    mu: float = None
 
 
-def tsinfer_workload(name, num_queries):
-    path = f"tgp_chr20p/{name}"
-    return Workload(name, f"{path}/reference.trees", path, num_queries)
+GENETIC_MAP = "genetic_maps/chr20_HapMapII_GRCh38.txt"
+
+
+def tsinfer_workload(name, num_queries, directory=None, **kwargs):
+    path = f"tgp_chr20p/{name if directory is None else directory}"
+    return Workload(name, f"{path}/reference.trees", path, num_queries, **kwargs)
 
 
 def sc2ts_workload(name, num_samples, num_queries, mismatch_threshold=None):
@@ -78,6 +88,10 @@ def sc2ts_workload(name, num_samples, num_queries, mismatch_threshold=None):
     tsinfer_workload("ancestors_g1484", 50),
     tsinfer_workload("ancestors_g1485", 100),
     tsinfer_workload("samples", 6),
+    tsinfer_workload("samples_map", 6, "samples", genetic_map=GENETIC_MAP),
+    tsinfer_workload(
+        "samples_map_mu1e-8", 6, "samples", genetic_map=GENETIC_MAP, mu=1e-8
+    ),
     sc2ts_workload("sc2ts_n100k_k1", 100000, 10, mismatch_threshold=1),
     sc2ts_workload("sc2ts_n100k", 100000, 10),
     sc2ts_workload("sc2ts_n300k", 300000, 4),
@@ -107,6 +121,12 @@ def load_workload(data_dir, workload):
         query_ids = SC2TS_QUERIES[: workload.num_queries]
     else:
         model = queries.load_model(data_dir / workload.query_dir)
+        if workload.genetic_map is not None:
+            model["rho"] = tsls.model.recombination_probabilities(
+                data_dir / workload.genetic_map, ts.sites_position
+            )
+        if workload.mu is not None:
+            model["mu"] = workload.mu
         query_ids = list(df.query_id[: workload.num_queries])
     return ts, H, query_ids, model
 
@@ -130,8 +150,8 @@ def time_workload(data_dir, workload, repeats):
     n = ts.num_nodes if model["weight_by_n"] else 1
     num_alleles = engines.site_num_alleles(ts, model["alleles"])
     rates = {
-        "rho": np.full(ts.num_sites, model["rho"]),
-        "mu": np.full(ts.num_sites, model["mu"]),
+        "rho": np.broadcast_to(model["rho"], ts.num_sites),
+        "mu": np.broadcast_to(model["mu"], ts.num_sites),
     }
     best = {name: np.full(len(query_ids), np.inf) for name in ENGINES}
     max_memory = {name: 0 for name in ENGINES}

benchmarking/profile_samples.py

@@ -0,0 +1,79 @@
+"""
+Run tsls Viterbi on chr20p sample queries, for profiling and step counts.
+
+The reference is the samples workload's reference.trees with the vestigial
+root edges from node 0 removed (as engines.TslsEngine does), and the model
+is tsinfer's: r = 0.01 or per-site probabilities from a genetic map, with
+weight_by_n=True (tsls's model). Prints the time and the engine's
+statistics for each query, and the step counters if the extension was
+built with TSLS_STEP_COUNTERS.
+
+    uv run python profile_samples.py --map data/genetic_maps/chr20_HapMapII_GRCh38.txt
+"""
+
+import hashlib
+import json
+import time
+
+import click
+import numpy as np
+import tskit
+import tsls
+
+import queries
+
+WORKLOAD = "data/tgp_chr20p/samples"
+
+
+def load_reference(path):
+    ts = tskit.load(path)
+    tables = ts.dump_tables()
+    tables.edges.keep_rows(tables.edges.parent != 0)
+    return tables.tree_sequence()
+
+
+@click.command()
+@click.option("--map", "genetic_map", default=None, help="HapMap genetic map")
+@click.option("--mu", type=float, default=1e-20)
+@click.option("--queries", "query_ids", default="0,40,80,120,160,199")
+@click.option("--missing", is_flag=True, help="Use all-missing haplotypes")
+@click.option("--repeats", type=int, default=1)
+def main(genetic_map, mu, query_ids, missing, repeats):
+    ts = load_reference(f"{WORKLOAD}/reference.trees")
+    H, positions, _ = queries.load_queries(WORKLOAD)
+    m = ts.num_sites
+    if genetic_map is None:
+        r = np.full(m, 0.01)
+    else:
+        r = tsls.model.recombination_probabilities(genetic_map, positions)
+    reference = tsls.TreeSequence(ts)
+    for qid in [int(q) for q in query_ids.split(",")]:
+        h = H[qid].astype(np.int8)
+        start, end = queries.query_window(H, qid)
+        if missing:
+            h[start:end] = -1
+        for _ in range(repeats):
+            hmm = tsls.Hmm(reference, r, np.full(m, mu))
+            matched = np.zeros(m, dtype=np.int8)
+            before = time.perf_counter()
+            left, right, parent = hmm.find_path(h, start, end, matched)
+            seconds = time.perf_counter() - before
+            digest = hashlib.md5()
+            for array in [left, right, parent, matched]:
+                digest.update(np.ascontiguousarray(array).tobytes())
+            row = {
+                "query": qid,
+                "time": round(seconds, 4),
+                "segments": len(left),
+                "path": digest.hexdigest()[:12],
+                "mean_values": round(hmm.mean_num_values, 1),
+                "traceback": round(hmm.mean_traceback_size, 1),
+            }
+            counts = getattr(hmm, "step_counts", None)
+            if counts is not None:
+                row.update(counts)
+            print(json.dumps(row), flush=True)
+
+
+if __name__ == "__main__":
+    main()