Initial tsls engine

.github/workflows/tests.yml

@@ -17,7 +17,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        package: [tsinfer-hmm, tskit-hmm]
+        package: [tsinfer-hmm, tskit-hmm, tsls]
     defaults:
       run:
         working-directory: ${{ matrix.package }}
@@ -44,7 +44,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        package: [tsinfer-hmm, tskit-hmm]
+        package: [tsinfer-hmm, tskit-hmm, tsls]
         python: ["3.11", "3.13"]
         os: [ubuntu-24.04, macos-latest]
     defaults:

README.md

@@ -21,6 +21,26 @@ Both build against the tskit C library in the `git-submodules/tskit`
 submodule, so clone with `--recurse-submodules` or run
 `git submodule update --init`.
 
+[`tsls/`](tsls) is the merged engine (Python package `tsls`, extension
+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.
+- 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 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.
+
+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 git history records exactly what was copied from where. Each
 directory was first added verbatim, then changed in separate commits to
 make it build on its own.
@@ -30,11 +50,12 @@ make it build on its own.
 tsls is distributed under the MIT License. tsinfer is licensed under
 GPL-3.0-or-later; the tsinfer code copied here was relicensed to MIT by
 its author. See [tsinfer-hmm/LICENSING.md](tsinfer-hmm/LICENSING.md) for
-details. The tskit code is MIT licensed.
+details. The merged engine in `tsls/` is derived from tsinfer-hmm and
+is covered by the same relicensing. The tskit code is MIT licensed.
 
 ## Development
 
-In either `tsinfer-hmm/` or `tskit-hmm/`:
+In `tsls/`, `tsinfer-hmm/` or `tskit-hmm/`:
 
 ```
 # Python tests

benchmarking/engines.py

@@ -30,6 +30,7 @@
 import numpy as np
 import tsinfer_hmm
 import tskit_hmm
+import tsls
 
 ALLELES = ("A", "C", "G", "T")
 MISSING = -1
@@ -202,4 +203,55 @@ def optimum_log_likelihood(self, normalisation_factor):
         return float(np.sum(np.log(normalisation_factor)) + np.log(n) - m * np.log(stay))
 
 
-ENGINES = {engine.name: engine for engine in [TsinferEngine, TskitEngine]}
+class TslsEngine(TsinferEngine):
+    """
+    The merged tsls engine, with the same model and options as tsinfer. In
+    tsls every node in the tree is a state, so we remove the vestigial root
+    edge 0 -> 1 added by prepare.py to make node 0 isolated, as in sc2ts.
+    """
+
+    name = "tsls"
+
+    def __init__(self, ts, *, mu, rho, mismatch_threshold=None):
+        self.num_sites = ts.num_sites
+        self.mu = mu
+        self.rho = rho
+        self.likelihood_threshold = None
+        if mismatch_threshold is not None:
+            self.likelihood_threshold = mu**mismatch_threshold * 0.99
+        self.recombination = np.full(ts.num_sites, rho)
+        self.mismatch = np.full(ts.num_sites, mu)
+        tables = ts.dump_tables()
+        tables.edges.keep_rows(tables.edges.parent != 0)
+        self.matcher_indexes = tsls.MatcherIndexes(
+            tables.tree_sequence(), alleles=ALLELES
+        )
+
+    def match(self, h):
+        matcher = tsls.AncestorMatcher(
+            self.matcher_indexes,
+            self.recombination,
+            self.mismatch,
+            likelihood_threshold=self.likelihood_threshold,
+            weight_by_n=False,
+        )
+        matched = np.zeros(self.num_sites, dtype=np.int8)
+        left, right, parent = matcher.find_path(
+            h.astype(np.int8), 0, self.num_sites, matched
+        )
+        stats = {
+            "total_memory": matcher.total_memory,
+            "mean_traceback_size": matcher.mean_traceback_size,
+        }
+        return (
+            left[::-1].astype(np.float64),
+            right[::-1].astype(np.float64),
+            parent[::-1].copy(),
+            matched,
+            stats,
+        )
+
+
+ENGINES = {
+    engine.name: engine for engine in [TsinferEngine, TskitEngine, TslsEngine]
+}

benchmarking/pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "tszip",
     "tsinfer-hmm",
     "tskit-hmm",
+    "tsls",
 ]
 
 [dependency-groups]
@@ -26,6 +27,7 @@ package = false
 [tool.uv.sources]
 tsinfer-hmm = { path = "../tsinfer-hmm", editable = true }
 tskit-hmm = { path = "../tskit-hmm", editable = true }
+tsls = { path = "../tsls", editable = true }
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

benchmarking/run.py

@@ -148,7 +148,7 @@ def main(
                 continue
             before = time.perf_counter()
             options = {}
-            if name == "tsinfer":
+            if name in ("tsinfer", "tsls"):
                 options["mismatch_threshold"] = mismatch_threshold
             engine = engines.ENGINES[name](ts, mu=mu, rho=rho, **options)
             setup_time = time.perf_counter() - before