Add skip_empty to Tree class and trees iterator

Author: hyanwong

python/tests/test_highlevel.py

@@ -1722,6 +1722,21 @@ def test_trees_interface(self):
             assert t.get_num_tracked_samples(0) == 0
             assert list(t.samples(0)) == [0]
 
+    def test_trees_skip_empty(self):
+        ts = tskit.Tree.generate_balanced(10).tree_sequence
+        ts = ts.delete_intervals([[0, 0.2], [0.7, 0.9]])
+        assert ts.num_trees == 4
+        for i, tree in enumerate(ts.trees()):
+            if i % 2 == 0:
+                assert tree.is_empty()
+            else:
+                assert not tree.is_empty()
+        assert i == 3
+
+        for _i, tree in enumerate(ts.trees(skip_empty=True)):
+            assert not tree.is_empty()
+        assert _i == 1
+
     @pytest.mark.parametrize("ts", get_example_tree_sequences())
     def test_get_pairwise_diversity(self, ts):
         with pytest.raises(ValueError, match="at least one element"):
@@ -4453,7 +4468,7 @@ class TestSeekDirection:
     def ts(self):
         return tsutil.all_trees_ts(3)
 
-    def setup(self):
+    def setup_trees(self):
         ts = self.ts()
         t1 = tskit.Tree(ts)
         t2 = tskit.Tree(ts)
@@ -4466,22 +4481,22 @@ def setup(self):
     def test_index_from_different_directions(self, index):
         # Check that we get different orderings of the children arrays
         # for all trees when we go in different directions.
-        t1, t2 = self.setup()
+        t1, t2 = self.setup_trees()
         while t1.index != index:
             t1.next()
         while t2.index != index:
             t2.prev()
         assert_same_tree_different_order(t1, t2)
 
     def test_seek_0_from_null(self):
-        t1, t2 = self.setup()
+        t1, t2 = self.setup_trees()
         t1.first()
         t2.seek(0)
         assert_trees_identical(t1, t2)
 
     @pytest.mark.parametrize("index", range(3))
     def test_seek_next_tree(self, index):
-        t1, t2 = self.setup()
+        t1, t2 = self.setup_trees()
         while t1.index != index:
             t1.next()
             t2.next()
@@ -4491,7 +4506,7 @@ def test_seek_next_tree(self, index):
 
     @pytest.mark.parametrize("index", [3, 2, 1])
     def test_seek_prev_tree(self, index):
-        t1, t2 = self.setup()
+        t1, t2 = self.setup_trees()
         while t1.index != index:
             t1.prev()
             t2.prev()
@@ -4500,44 +4515,44 @@ def test_seek_prev_tree(self, index):
         assert_trees_identical(t1, t2)
 
     def test_seek_1_from_0(self):
-        t1, t2 = self.setup()
+        t1, t2 = self.setup_trees()
         t1.first()
         t1.next()
         t2.first()
         t2.seek(1)
         assert_trees_identical(t1, t2)
 
     def test_seek_1_5_from_0(self):
-        t1, t2 = self.setup()
+        t1, t2 = self.setup_trees()
         t1.first()
         t1.next()
         t2.first()
         t2.seek(1.5)
         assert_trees_identical(t1, t2)
 
     def test_seek_1_5_from_1(self):
-        t1, t2 = self.setup()
+        t1, t2 = self.setup_trees()
         for _ in range(2):
             t1.next()
             t2.next()
         t2.seek(1.5)
         assert_trees_identical(t1, t2)
 
     def test_seek_3_from_null(self):
-        t1, t2 = self.setup()
+        t1, t2 = self.setup_trees()
         t1.last()
         t2.seek(3)
         assert_trees_identical(t1, t2)
 
     def test_seek_3_from_0(self):
-        t1, t2 = self.setup()
+        t1, t2 = self.setup_trees()
         t1.last()
         t2.first()
         t2.seek(3)
         assert_trees_identical(t1, t2)
 
     def test_seek_0_from_3(self):
-        t1, t2 = self.setup()
+        t1, t2 = self.setup_trees()
         t1.last()
         t1.first()
         t2.last()

python/tskit/trees.py

@@ -652,6 +652,7 @@ def __init__(
         tracked_samples=None,
         *,
         sample_lists=False,
+        skip_empty=None,
         root_threshold=1,
         sample_counts=None,
     ):
@@ -671,6 +672,7 @@ def __init__(
 
         self._tree_sequence = tree_sequence
         self._ll_tree = _tskit.Tree(tree_sequence.ll_tree_sequence, **kwargs)
+        self.skip_empty = skip_empty  # This is a python-only facility, not in _ll
         self._ll_tree.set_root_threshold(root_threshold)
         self._make_arrays()
 
@@ -3851,6 +3853,7 @@ def __init__(self, tree):
         self.tree = tree
         self.more_trees = True
         self.forward = True
+        self.no_skip_empty = True if tree.skip_empty is None else not tree.skip_empty
 
     def __iter__(self):
         return self
@@ -3860,15 +3863,19 @@ def __reversed__(self):
         return self
 
     def __next__(self):
-        if self.forward:
-            self.more_trees = self.more_trees and self.tree.next()
-        else:
-            self.more_trees = self.more_trees and self.tree.prev()
-        if not self.more_trees:
-            raise StopIteration()
+        while True:
+            if self.forward:
+                self.more_trees = self.more_trees and self.tree.next()
+            else:
+                self.more_trees = self.more_trees and self.tree.prev()
+            if not self.more_trees:
+                raise StopIteration()
+            if self.no_skip_empty or not self.tree.is_empty():
+                break
         return self.tree
 
     def __len__(self):
+        # NB: this can return a longer length than the iterator if we skip empty trees
         return self.tree.tree_sequence.num_trees
 
 
@@ -4945,6 +4952,7 @@ def trees(
         tracked_samples=None,
         *,
         sample_lists=False,
+        skip_empty=None,
         root_threshold=1,
         sample_counts=None,
         tracked_leaves=None,
@@ -4972,6 +4980,9 @@ def trees(
         :param bool sample_lists: If True, provide more efficient access
             to the samples beneath a given node using the
             :meth:`Tree.samples` method.
+        :param bool skip_empty: If True, skip trees that are
+            :meth:`empty<Tree.is_empty>`, commonly found at the start and end of
+            a tree sequence. Default: ``None`` treated as ``False``.
         :param int root_threshold: The minimum number of samples that a node
             must be ancestral to for it to be in the list of roots. By default
             this is 1, so that isolated samples (representing missing data)
@@ -4996,6 +5007,7 @@ def trees(
             self,
             tracked_samples=tracked_samples,
             sample_lists=sample_lists,
+            skip_empty=skip_empty,
             root_threshold=root_threshold,
             sample_counts=sample_counts,
         )