Move add_vestigial_root to new arg_ops module

Create tsinfer/arg_ops.py for ARG topology operations and move
add_vestigial_root from matching.py. Create tests/test_arg_ops.py
with the corresponding tests. Update all call sites in matching.py,
test_python_c.py, and test_lshmm.py.

Author: jeromekelleher

tests/test_arg_ops.py

@@ -0,0 +1,19 @@
+import pytest
+import tskit
+
+from tsinfer import arg_ops
+
+
+class TestAddVestigialRoot:
+    def test_non_discrete_genome(self):
+        tables = tskit.TableCollection(sequence_length=1.5)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        ts = tables.tree_sequence()
+        with pytest.raises(ValueError, match="discrete genome"):
+            arg_ops.add_vestigial_root(ts)
+
+    def test_empty_tree_sequence(self):
+        tables = tskit.TableCollection(sequence_length=1)
+        ts = tables.tree_sequence()
+        with pytest.raises(ValueError, match="Emtpy trees"):
+            arg_ops.add_vestigial_root(ts)

tests/test_lshmm.py

@@ -14,7 +14,7 @@
 import tskit
 
 import _tsinfer
-from tsinfer import matching
+from tsinfer import arg_ops, matching
 
 
 @dataclasses.dataclass
@@ -78,7 +78,7 @@ class MatcherIndexes:
     def __init__(self, in_tables, *, vestigial_root=True, num_alleles=None):
         ts = in_tables.tree_sequence()
         if vestigial_root:
-            ts = matching.add_vestigial_root(ts)
+            ts = arg_ops.add_vestigial_root(ts)
         tables = ts.dump_tables()
 
         self.sequence_length = tables.sequence_length

tests/test_matching.py

@@ -23,7 +23,6 @@
 from __future__ import annotations
 
 import numpy as np
-import pytest
 import tskit
 
 from tsinfer import grouping, matching, vcz
@@ -761,26 +760,6 @@ def test_metadata_survives_multiple_cycles(self):
         assert ts2.metadata["sequence_intervals"] == [[10, 51]]
 
 
-# ---------------------------------------------------------------------------
-# TestAddVestigialRoot
-# ---------------------------------------------------------------------------
-
-
-class TestAddVestigialRoot:
-    def test_non_discrete_genome(self):
-        tables = tskit.TableCollection(sequence_length=1.5)
-        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
-        ts = tables.tree_sequence()
-        with pytest.raises(ValueError, match="discrete genome"):
-            matching.add_vestigial_root(ts)
-
-    def test_empty_tree_sequence(self):
-        tables = tskit.TableCollection(sequence_length=1)
-        ts = tables.tree_sequence()
-        with pytest.raises(ValueError, match="Emtpy trees"):
-            matching.add_vestigial_root(ts)
-
-
 # ---------------------------------------------------------------------------
 # TestAncestorMatcherWrapper
 # ---------------------------------------------------------------------------

tests/test_python_c.py

@@ -28,7 +28,7 @@
 import tskit
 
 import _tsinfer
-from tsinfer import matching
+from tsinfer import arg_ops
 
 IS_WINDOWS = sys.platform == "win32"
 
@@ -155,7 +155,7 @@ def make_matcher_indexes_and_matcher(num_samples=4):
     tables.sites.add_row(position=1, ancestral_state="A")
     tables.mutations.add_row(site=0, node=1, derived_state="T")
     ts = tables.tree_sequence()
-    ts = matching.add_vestigial_root(ts)
+    ts = arg_ops.add_vestigial_root(ts)
     ll_tables = _tsinfer.LightweightTableCollection(ts.sequence_length)
     ll_tables.fromdict(ts.dump_tables().asdict())
     mi = _tsinfer.MatcherIndexes(ll_tables)
@@ -267,7 +267,7 @@ def test_find_path_match_impossible(self):
         tables.sites.add_row(position=1, ancestral_state="A")
         # No mutations: all nodes carry allele 0
         ts = tables.tree_sequence()
-        ts = matching.add_vestigial_root(ts)
+        ts = arg_ops.add_vestigial_root(ts)
         ll_tables = _tsinfer.LightweightTableCollection(ts.sequence_length)
         ll_tables.fromdict(ts.dump_tables().asdict())
         mi = _tsinfer.MatcherIndexes(ll_tables)
@@ -299,7 +299,7 @@ def test_get_traceback_bad_site(self):
 class TestMatcherIndexes:
     def test_single_tree(self):
         ts = tskit.Tree.generate_balanced(4).tree_sequence
-        ts = matching.add_vestigial_root(ts)
+        ts = arg_ops.add_vestigial_root(ts)
         tables = ts.dump_tables()
         ll_tables = _tsinfer.LightweightTableCollection(tables.sequence_length)
         ll_tables.fromdict(tables.asdict())
@@ -323,7 +323,7 @@ def test_num_alleles(self):
 
     def test_print_state(self, tmpdir):
         ts = tskit.Tree.generate_balanced(4).tree_sequence
-        ts = matching.add_vestigial_root(ts)
+        ts = arg_ops.add_vestigial_root(ts)
         tables = ts.dump_tables()
         ll_tables = _tsinfer.LightweightTableCollection(tables.sequence_length)
         ll_tables.fromdict(tables.asdict())
@@ -341,7 +341,7 @@ def test_print_state(self, tmpdir):
 
     def test_print_state_bad_file(self):
         ts = tskit.Tree.generate_balanced(4).tree_sequence
-        ts = matching.add_vestigial_root(ts)
+        ts = arg_ops.add_vestigial_root(ts)
         tables = ts.dump_tables()
         ll_tables = _tsinfer.LightweightTableCollection(tables.sequence_length)
         ll_tables.fromdict(tables.asdict())

tsinfer/arg_ops.py

@@ -0,0 +1,61 @@
+#
+# Copyright (C) 2018-2026 University of Oxford
+#
+# This file is part of tsinfer.
+#
+# tsinfer is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# tsinfer is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with tsinfer.  If not, see <http://www.gnu.org/licenses/>.
+#
+"""
+Operations on ARG (Ancestral Recombination Graph) topology.
+"""
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def add_vestigial_root(ts):
+    """
+    Adds the nodes and edges required by tsinfer to the specified tree sequence
+    and returns it.
+    """
+    if not ts.discrete_genome:
+        raise ValueError("Only discrete genome coords supported")
+    if ts.num_nodes == 0:
+        raise ValueError("Emtpy trees not supported")
+
+    base_tables = ts.dump_tables()
+    tables = base_tables.copy()
+    tables.nodes.clear()
+    t = max(ts.nodes_time)
+    tables.nodes.add_row(time=t + 1)
+    num_additonal_nodes = 1
+    tables.mutations.node += num_additonal_nodes
+    tables.edges.child += num_additonal_nodes
+    tables.edges.parent += num_additonal_nodes
+    for node in base_tables.nodes:
+        tables.nodes.append(node)
+    if ts.num_edges > 0:
+        for tree in ts.trees():
+            # if tree.num_roots > 1:
+            #     print(ts.draw_text())
+            root = tree.root + num_additonal_nodes
+            tables.edges.add_row(
+                tree.interval.left, tree.interval.right, parent=0, child=root
+            )
+        tables.edges.squash()
+        # FIXME probably don't need to sort here most of the time, or at least
+        # we can just sort almost the end of the table.
+        tables.sort()
+    return tables.tree_sequence()

tsinfer/matching.py

@@ -33,7 +33,7 @@
 
 import _tsinfer
 
-from . import grouping, vcz
+from . import arg_ops, grouping, vcz
 
 logger = logging.getLogger(__name__)
 
@@ -380,48 +380,12 @@ def extend_ts(
     return result_ts
 
 
-def add_vestigial_root(ts):
-    """
-    Adds the nodes and edges required by tsinfer to the specified tree sequence
-    and returns it.
-    """
-    if not ts.discrete_genome:
-        raise ValueError("Only discrete genome coords supported")
-    if ts.num_nodes == 0:
-        raise ValueError("Emtpy trees not supported")
-
-    base_tables = ts.dump_tables()
-    tables = base_tables.copy()
-    tables.nodes.clear()
-    t = max(ts.nodes_time)
-    tables.nodes.add_row(time=t + 1)
-    num_additonal_nodes = 1
-    tables.mutations.node += num_additonal_nodes
-    tables.edges.child += num_additonal_nodes
-    tables.edges.parent += num_additonal_nodes
-    for node in base_tables.nodes:
-        tables.nodes.append(node)
-    if ts.num_edges > 0:
-        for tree in ts.trees():
-            # if tree.num_roots > 1:
-            #     print(ts.draw_text())
-            root = tree.root + num_additonal_nodes
-            tables.edges.add_row(
-                tree.interval.left, tree.interval.right, parent=0, child=root
-            )
-        tables.edges.squash()
-        # FIXME probably don't need to sort here most of the time, or at least we
-        # can just sort almost the end of the table.
-        tables.sort()
-    return tables.tree_sequence()
-
-
 class MatcherIndexes(_tsinfer.MatcherIndexes):
     """Wrapper around the C MatcherIndexes, built from a tree sequence."""
 
     def __init__(self, ts, *, vestigial_root=True, num_alleles=None):
         if vestigial_root:
-            ts = add_vestigial_root(ts)
+            ts = arg_ops.add_vestigial_root(ts)
         tables = ts.dump_tables()
         ll_tables = _tsinfer.LightweightTableCollection(tables.sequence_length)
         ll_tables.fromdict(tables.asdict())