Add ts.coiterate method

Fixes #1021

Author: hyanwong

python/tests/test_topology.py

@@ -4785,6 +4785,86 @@ def test_partial_overlap_contradictory_children(self):
             tskit.load_text(nodes=nodes, edges=edges, strict=False)
 
 
+class TestCoiteration:
+    """
+    Test ability to iterate over multiple (currently 2) tree sequences simultaneously
+    """
+
+    def test_identical_ts(self):
+        ts = msprime.simulate(4, recombination_rate=1, random_seed=123)
+        assert ts.num_trees > 1
+        total_iterations = 0
+        for _, t1, t2 in ts.coiterate(ts):
+            total_iterations += 1
+            assert t1 == t2
+        assert ts.num_trees == total_iterations
+
+    def test_intervals(self):
+        ts1 = msprime.simulate(4, recombination_rate=1, random_seed=1)
+        assert ts1.num_trees > 1
+        one_tree_ts = msprime.simulate(5, random_seed=2)
+        multi_tree_ts = msprime.simulate(5, recombination_rate=1, random_seed=2)
+        assert multi_tree_ts.num_trees > 1
+        for ts2 in (one_tree_ts, multi_tree_ts):
+            bp1 = set(ts1.breakpoints())
+            bp2 = set(ts2.breakpoints())
+            assert bp1 != bp2
+            breaks = {0}
+            for interval, t1, t2 in ts1.coiterate(ts2):
+                breaks.add(interval.right)
+                assert t1.tree_sequence == ts1
+                assert t2.tree_sequence == ts2
+            assert breaks == bp1 | bp2
+
+    def test_simple_ts(self):
+        nodes = """\
+        id      is_sample   time
+        0       1           0
+        1       1           0
+        2       1           0
+        3       0           1
+        4       0           2
+        """
+        edges1 = """\
+        left    right   parent  child
+        0       0.2       3       0,1
+        0       0.2       4       2,3
+        0.2     1         3       2,1
+        0.2     1         4       0,3
+        """
+        edges2 = """\
+        left    right   parent  child
+        0       0.8       3       2,1
+        0       0.8       4       0,3
+        0.8     1         3       0,1
+        0.8     1         4       2,3
+        """
+        ts1 = tskit.load_text(io.StringIO(nodes), io.StringIO(edges1), strict=False)
+        ts2 = tskit.load_text(io.StringIO(nodes), io.StringIO(edges2), strict=False)
+        coiterator = ts1.coiterate(ts2)
+        interval, tree1, tree2 = next(coiterator)
+        assert interval.left == 0
+        assert interval.right == 0.2
+        assert tree1 == ts1.at_index(0)
+        assert tree2 == ts2.at_index(0)
+        interval, tree1, tree2 = next(coiterator)
+        assert interval.left == 0.2
+        assert interval.right == 0.8
+        assert tree1 == ts1.at_index(1)
+        assert tree2 == ts2.at_index(0)
+        interval, tree1, tree2 = next(coiterator)
+        assert interval.left == 0.8
+        assert interval.right == 1
+        assert tree1 == ts1.at_index(1)
+        assert tree2 == ts2.at_index(1)
+
+    def test_nonequal_lengths(self):
+        ts1 = msprime.simulate(4, random_seed=1, length=2)
+        ts2 = msprime.simulate(4, random_seed=1)
+        with pytest.raises(ValueError):
+            next(ts1.coiterate(ts2))
+
+
 class SimplifyTestBase:
     """
     Base class for simplify tests.

python/tests/tsutil.py

@@ -1282,27 +1282,3 @@ def genealogical_nearest_neighbours(ts, focal, reference_sets):
     L[L == 0] = 1
     A /= L.reshape((len(focal), 1))
     return A
-
-
-def coiterate(ts1, ts2, **kwargs):
-    """
-    Returns an iterator over the pairs of trees for each distinct
-    interval in the specified pair of tree sequences.
-    """
-    if ts1.sequence_length != ts2.sequence_length:
-        raise ValueError("Tree sequences must be equal length.")
-    L = ts1.sequence_length
-    trees1 = ts1.trees(**kwargs)
-    trees2 = ts2.trees(**kwargs)
-    tree1 = next(trees1)
-    tree2 = next(trees2)
-    right = 0
-    while right != L:
-        left = right
-        right = min(tree1.interval[1], tree2.interval[1])
-        yield (left, right), tree1, tree2
-        # Advance
-        if tree1.interval[1] == right:
-            tree1 = next(trees1, None)
-        if tree2.interval[1] == right:
-            tree2 = next(trees2, None)

python/tskit/trees.py

@@ -3850,6 +3850,43 @@ def trees(
         )
         return TreeIterator(tree)
 
+    def coiterate(self, other, **kwargs):
+        """
+        Returns an iterator over the pairs of trees for each distinct
+        interval in the specified pair of tree sequences.
+
+        :param bool other: The other tree sequence from which to take trees. The
+            sequence length should be the same as the current tree sequence.
+        :param \\**kwargs: Further named arguments that will be passed to the
+            :meth:`.trees` method when constructing the returned trees.
+
+        :return: An iterator returning successive tuples of the form
+            ``(interval, tree_self, tree_other)``. For example, the first item returned
+            will consist of an tuple of the intail interval, the first tree of the
+            current tree sequence, and the first tree of the ``other`` tree sequence;
+            the ``.left`` attribute of the initial interval will be 0 and the ``.right``
+            attribute will be the smallest non-zero breakpoint of the 2 tree sequences.
+        :rtype: iter(:class:`Interval`, :class:`Tree`, :class:`Tree`)
+
+        """
+        if self.sequence_length != other.sequence_length:
+            raise ValueError("Tree sequences must be equal length.")
+        L = self.sequence_length
+        trees1 = self.trees(**kwargs)
+        trees2 = other.trees(**kwargs)
+        tree1 = next(trees1)
+        tree2 = next(trees2)
+        right = 0
+        while right != L:
+            left = right
+            right = min(tree1.interval[1], tree2.interval[1])
+            yield Interval(left, right), tree1, tree2
+            # Advance
+            if tree1.interval[1] == right:
+                tree1 = next(trees1, None)
+            if tree2.interval[1] == right:
+                tree2 = next(trees2, None)
+
     def haplotypes(
         self,
         *,