Python version of virtual roots in tree arrays.

Includes tests to document differences in semantics to legacy version.

Author: jeromekelleher

python/tests/__init__.py

@@ -297,7 +297,7 @@ def edge_diffs(self):
             left = right
 
     def trees(self):
-        rtt = tsutil.RootThresholdTree(self._tree_sequence)
+        rtt = tsutil.LegacyRootThresholdTree(self._tree_sequence)
         pt = PythonTree(self._tree_sequence.get_num_nodes())
         pt.index = 0
         for left, right in rtt.iterate():

python/tests/test_topology.py

@@ -1,6 +1,6 @@
 # MIT License
 #
-# Copyright (c) 2018-2020 Tskit Developers
+# Copyright (c) 2018-2021 Tskit Developers
 # Copyright (c) 2016-2017 University of Oxford
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -323,6 +323,12 @@ def test_multiroot_tree(self):
         ts = tsutil.decapitate(ts, ts.num_edges // 2)
         self.verify(ts)
 
+    def test_all_missing_data(self):
+        tables = tskit.TableCollection(1)
+        for _ in range(10):
+            tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        self.verify(tables.tree_sequence())
+
 
 class TestKCMetric(unittest.TestCase):
     """
@@ -6605,6 +6611,11 @@ def verify(self, ts):
                             break
                         index = tree1.next_sample[index]
                 assert samples1 == samples2
+            np.testing.assert_array_equal(tree1.parent, tree2.parent_array)
+            np.testing.assert_array_equal(tree1.left_child, tree2.left_child_array)
+            np.testing.assert_array_equal(tree1.right_child, tree2.right_child_array)
+            # We don't compare the sib arrays because these aren't maintained properly
+            # for roots in the Python implementation
         assert right == ts.sequence_length
 
 
@@ -6615,7 +6626,7 @@ class TestOneSampleRoot(ExampleTopologyMixin):
     """
 
     def verify(self, ts):
-        tree1 = tsutil.RootThresholdTree(ts, root_threshold=1)
+        tree1 = tsutil.LegacyRootThresholdTree(ts, root_threshold=1)
         tree2 = tskit.Tree(ts)
         tree2.first()
         for interval in tree1.iterate():
@@ -6629,35 +6640,102 @@ def verify(self, ts):
                     u = tree2.parent(u)
                 roots.add(path_end)
             assert set(tree1.roots()) == roots
+            np.testing.assert_array_equal(tree1.parent, tree2.parent_array)
+            np.testing.assert_array_equal(tree1.left_child, tree2.left_child_array)
+            np.testing.assert_array_equal(tree1.right_child, tree2.right_child_array)
+            np.testing.assert_array_equal(tree1.left_sib, tree2.left_sib_array)
+            np.testing.assert_array_equal(tree1.right_sib, tree2.right_sib_array)
             tree2.next()
         assert tree2.index == -1
 
 
-class TestKSamplesRoot(ExampleTopologyMixin):
+class RootThreshold(ExampleTopologyMixin):
     """
     Tests for the root criteria of subtending at least k samples.
     """
 
+    @pytest.mark.skip("all missing broken in lib: #1706")
+    def test_all_missing_data(self):
+        pass
+
     def verify(self, ts):
-        for k in range(1, 5):
-            tree1 = tsutil.RootThresholdTree(ts, root_threshold=k)
-            tree2 = tskit.Tree(ts, root_threshold=k)
-            tree2.first()
-            for interval in tree1.iterate():
-                assert interval == tree2.interval
-                # Definition here is the set unique path ends from samples
-                # that subtend at least k samples
-                roots = set()
-                for u in ts.samples():
-                    while u != tskit.NULL:
-                        path_end = u
-                        u = tree2.parent(u)
-                    if tree2.num_samples(path_end) >= k:
-                        roots.add(path_end)
-                assert set(tree1.roots()) == roots
-                assert tree1.roots() == tree2.roots
-                tree2.next()
-            assert tree2.index == -1
+        k = self.root_threshold
+        trees_py = tsutil.algorithm_R(ts, root_threshold=k)
+        tree_lib = tskit.Tree(ts, root_threshold=k)
+        tree_lib.first()
+        tree_leg = tsutil.LegacyRootThresholdTree(ts, root_threshold=k)
+        for (interval_py, tree_py), interval_leg in itertools.zip_longest(
+            trees_py, tree_leg.iterate()
+        ):
+            assert interval_py == tree_lib.interval
+            assert interval_leg == tree_lib.interval
+
+            # Definition here is the set unique path ends from samples
+            # that subtend at least k samples
+            roots = set()
+            for u in ts.samples():
+                while u != tskit.NULL:
+                    path_end = u
+                    u = tree_lib.parent(u)
+                if tree_lib.num_samples(path_end) >= k:
+                    roots.add(path_end)
+            assert set(tree_py.roots()) == roots
+            assert set(tree_lib.roots) == roots
+            assert tree_leg.roots() == tree_lib.roots
+            assert len(tree_py.roots()) == tree_lib.num_roots
+            assert len(tree_leg.roots()) == tree_lib.num_roots
+
+            # The legacy class has identical behaviour to the lib version
+            assert tree_leg.left_root == tree_lib.left_root
+            np.testing.assert_array_equal(tree_leg.parent, tree_lib.parent_array)
+            np.testing.assert_array_equal(
+                tree_leg.left_child, tree_lib.left_child_array
+            )
+            np.testing.assert_array_equal(
+                tree_leg.right_child, tree_lib.right_child_array
+            )
+            np.testing.assert_array_equal(tree_leg.left_sib, tree_lib.left_sib_array)
+            np.testing.assert_array_equal(tree_leg.right_sib, tree_lib.right_sib_array)
+
+            # NOTE: the legacy left_root value is *not* necessarily the same as the
+            # new left_root.
+            # assert tree_leg.left_root == tree_py.left_child[-1]
+
+            # The virtual root version is identical except for the extra node and
+            # the details of the sib arrays.
+            np.testing.assert_array_equal(tree_py.parent[:-1], tree_leg.parent)
+            np.testing.assert_array_equal(tree_py.left_child[:-1], tree_leg.left_child)
+            np.testing.assert_array_equal(
+                tree_py.right_child[:-1], tree_leg.right_child
+            )
+            # The sib arrays are identical except for root nodes.
+            for u in range(ts.num_nodes):
+                if u not in roots:
+                    assert tree_py.left_sib[u] == tree_leg.left_sib[u]
+                    assert tree_py.right_sib[u] == tree_leg.right_sib[u]
+
+            tree_lib.next()
+        assert tree_lib.index == -1
+
+
+class TestRootThreshold1(RootThreshold):
+    root_threshold = 1
+
+
+class TestRootThreshold2(RootThreshold):
+    root_threshold = 2
+
+
+class TestRootThreshold3(RootThreshold):
+    root_threshold = 3
+
+
+class TestRootThreshold4(RootThreshold):
+    root_threshold = 4
+
+
+class TestRootThreshold10(RootThreshold):
+    root_threshold = 10
 
 
 class TestSquashEdges:

python/tests/tsutil.py

@@ -1,6 +1,6 @@
 # MIT License
 #
-# Copyright (c) 2018-2019 Tskit Developers
+# Copyright (c) 2018-2021 Tskit Developers
 # Copyright (C) 2017 University of Oxford
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -1175,6 +1175,143 @@ def algorithm_T(ts):
         left = right
 
 
+class QuintuplyLinkedTree:
+    def __init__(self, n, root_threshold=1):
+        self.root_threshold = root_threshold
+        self.parent = np.zeros(n + 1, dtype=np.int32) - 1
+        self.left_child = np.zeros(n + 1, dtype=np.int32) - 1
+        self.right_child = np.zeros(n + 1, dtype=np.int32) - 1
+        self.left_sib = np.zeros(n + 1, dtype=np.int32) - 1
+        self.right_sib = np.zeros(n + 1, dtype=np.int32) - 1
+        self.num_samples = np.zeros(n + 1, dtype=np.int32)
+
+    def __str__(self):
+        s = "id\tparent\tlchild\trchild\tlsib\trsib\tnsamp\n"
+        for j in range(len(self.parent)):
+            s += (
+                f"{j}\t{self.parent[j]}\t"
+                f"{self.left_child[j]}\t{self.right_child[j]}\t"
+                f"{self.left_sib[j]}\t{self.right_sib[j]}\t"
+                f"{self.num_samples[j]}\n"
+            )
+        return s
+
+    def roots(self):
+        roots = []
+        u = self.left_child[-1]
+        while u != -1:
+            roots.append(u)
+            u = self.right_sib[u]
+        return roots
+
+    def remove_branch(self, p, c):
+        lsib = self.left_sib[c]
+        rsib = self.right_sib[c]
+        if lsib == -1:
+            self.left_child[p] = rsib
+        else:
+            self.right_sib[lsib] = rsib
+        if rsib == -1:
+            self.right_child[p] = lsib
+        else:
+            self.left_sib[rsib] = lsib
+        self.parent[c] = -1
+        self.left_sib[c] = -1
+        self.right_sib[c] = -1
+
+    def insert_branch(self, p, c):
+        assert self.parent[c] == -1, "contradictory edges"
+        self.parent[c] = p
+        u = self.right_child[p]
+        if u == -1:
+            self.left_child[p] = c
+            self.left_sib[c] = -1
+            self.right_sib[c] = -1
+        else:
+            self.right_sib[u] = c
+            self.left_sib[c] = u
+            self.right_sib[c] = -1
+        self.right_child[p] = c
+
+    def is_root(self, u):
+        return self.num_samples[u] >= self.root_threshold
+
+    # Note we cheat a bit here and use the -1 == last element semantics from Python.
+    # We could use self.insert_branch(N, root) and then set self.parent[root] = -1.
+    def insert_root(self, root):
+        self.insert_branch(-1, root)
+
+    def remove_root(self, root):
+        self.remove_branch(-1, root)
+
+    def remove_edge(self, edge):
+        self.remove_branch(edge.parent, edge.child)
+
+        u = edge.parent
+        while u != -1:
+            path_end = u
+            path_end_was_root = self.is_root(u)
+            self.num_samples[u] -= self.num_samples[edge.child]
+            u = self.parent[u]
+
+        if path_end_was_root and not self.is_root(path_end):
+            self.remove_root(path_end)
+        if self.is_root(edge.child):
+            self.insert_root(edge.child)
+
+    def insert_edge(self, edge):
+        u = edge.parent
+        while u != -1:
+            path_end = u
+            path_end_was_root = self.is_root(u)
+            self.num_samples[u] += self.num_samples[edge.child]
+            u = self.parent[u]
+
+        if self.is_root(edge.child):
+            self.remove_root(edge.child)
+        if self.is_root(path_end) and not path_end_was_root:
+            self.insert_root(path_end)
+
+        self.insert_branch(edge.parent, edge.child)
+
+
+def algorithm_R(ts, root_threshold=1):
+    """
+    Quintuply linked tree with root tracking.
+    """
+    sequence_length = ts.sequence_length
+    N = ts.num_nodes
+    M = ts.num_edges
+    tree = QuintuplyLinkedTree(N, root_threshold=root_threshold)
+    edges = list(ts.edges())
+    in_order = ts.tables.indexes.edge_insertion_order
+    out_order = ts.tables.indexes.edge_removal_order
+
+    # Initialise the tree
+    for u in ts.samples():
+        tree.num_samples[u] = 1
+        if tree.is_root(u):
+            tree.insert_root(u)
+
+    j = 0
+    k = 0
+    left = 0
+    while j < M or left < sequence_length:
+        while k < M and edges[out_order[k]].right == left:
+            tree.remove_edge(edges[out_order[k]])
+            k += 1
+        while j < M and edges[in_order[j]].left == left:
+            tree.insert_edge(edges[in_order[j]])
+            j += 1
+        right = sequence_length
+        if j < M:
+            right = min(right, edges[in_order[j]].left)
+        if k < M:
+            right = min(right, edges[out_order[k]].right)
+        yield (left, right), tree
+        left = right
+
+
 class SampleListTree:
     """
     Straightforward implementation of the quintuply linked tree for developing
@@ -1315,15 +1452,8 @@ def sample_lists(self):
         sequence_length = ts.sequence_length
         edges = list(ts.edges())
         M = len(edges)
-        time = [ts.node(edge.parent).time for edge in edges]
-        in_order = sorted(
-            range(M),
-            key=lambda j: (edges[j].left, time[j], edges[j].parent, edges[j].child),
-        )
-        out_order = sorted(
-            range(M),
-            key=lambda j: (edges[j].right, -time[j], -edges[j].parent, -edges[j].child),
-        )
+        in_order = ts.tables.indexes.edge_insertion_order
+        out_order = ts.tables.indexes.edge_removal_order
         j = 0
         k = 0
         left = 0
@@ -1348,10 +1478,12 @@ def sample_lists(self):
             left = right
 
 
-class RootThresholdTree:
+class LegacyRootThresholdTree:
     """
-    Straightforward implementation of the quintuply linked tree for developing
-    and testing the root_threshold feature.
+    Implementation of the quintuply linked tree with root tracking using the
+    pre C 1.0/Python 0.4.0 algorithm. We keep this version around to make sure
+    that we can be clear what the differences in the semantics of the new
+    and old versions are.
 
     NOTE: The interface is pretty awkward; it's not intended for anything other
     than testing.
@@ -1512,15 +1644,8 @@ def iterate(self):
         sequence_length = ts.sequence_length
         edges = list(ts.edges())
         M = len(edges)
-        time = [ts.node(edge.parent).time for edge in edges]
-        in_order = sorted(
-            range(M),
-            key=lambda j: (edges[j].left, time[j], edges[j].parent, edges[j].child),
-        )
-        out_order = sorted(
-            range(M),
-            key=lambda j: (edges[j].right, -time[j], -edges[j].parent, -edges[j].child),
-        )
+        in_order = ts.tables.indexes.edge_insertion_order
+        out_order = ts.tables.indexes.edge_removal_order
         j = 0
         k = 0
         left = 0