Work on traversal orders in Python

Author: jeromekelleher

python/tests/__init__.py

@@ -83,80 +83,6 @@ def children(self, u):
             v = self.right_sib[v]
         return ret
 
-    def _preorder_nodes(self, u, node_ist):
-        node_ist.append(u)
-        for c in self.children(u):
-            self._preorder_nodes(c, node_ist)
-
-    def _postorder_nodes(self, u, node_ist):
-        for c in self.children(u):
-            self._postorder_nodes(c, node_ist)
-        node_ist.append(u)
-
-    def _inorder_nodes(self, u, node_ist):
-        children = self.children(u)
-        if len(children) > 0:
-            mid = len(children) // 2
-            for v in children[:mid]:
-                self._inorder_nodes(v, node_ist)
-            node_ist.append(u)
-            for v in children[mid:]:
-                self._inorder_nodes(v, node_ist)
-        else:
-            node_ist.append(u)
-
-    def _levelorder_nodes(self, u, node_ist, level):
-        node_ist[level].append(u) if level < len(node_ist) else node_ist.append([u])
-        for c in self.children(u):
-            self._levelorder_nodes(c, node_ist, level + 1)
-
-    def _minlex_postorder_nodes(self, u, node_ist):
-        node_ist.extend(self._minlex_postorder_nodes_helper(u)[1])
-
-    def _minlex_postorder_nodes_helper(self, u):
-        """
-        For a given input ID u, this function returns a tuple whose first value
-        is the minimum leaf node ID under node u, and whose second value is
-        a list containing the minlex postorder for the subtree rooted at node u.
-        The first value is needed for sorting, and the second value is what
-        finally gets returned.
-        """
-        children = self.children(u)
-        if len(children) > 0:
-            children_return = [self._minlex_postorder_nodes_helper(c) for c in children]
-            # sorts by first value, which is the minimum leaf node ID
-            children_return.sort()
-            minlex_postorder = []
-            for _, child_minlex_postorder in children_return:
-                minlex_postorder.extend(child_minlex_postorder)
-            minlex_postorder.extend([u])
-            return (children_return[0][0], minlex_postorder)
-        else:
-            return (u, [u])
-
-    def nodes(self, root=None, order="preorder"):
-        roots = [root]
-        if root is None:
-            roots = self.roots
-        for u in roots:
-            node_list = []
-            if order == "preorder":
-                self._preorder_nodes(u, node_list)
-            elif order == "inorder":
-                self._inorder_nodes(u, node_list)
-            elif order == "postorder":
-                self._postorder_nodes(u, node_list)
-            elif order == "levelorder" or order == "breadthfirst":
-                # Returns nodes in their respective levels
-                # Nested list comprehension flattens node_list in order
-                self._levelorder_nodes(u, node_list, 0)
-                node_list = iter([i for level in node_list for i in level])
-            elif order == "minlex_postorder":
-                self._minlex_postorder_nodes(u, node_list)
-            else:
-                raise ValueError("order not supported")
-            yield from node_list
-
     def get_interval(self):
         return self.left, self.right
 

python/tests/test_highlevel.py

@@ -58,6 +58,103 @@
 from tskit import UNKNOWN_TIME
 
 
+def traversal_preorder(tree, root=None):
+    roots = tree.roots if root is None else [root]
+    for node in roots:
+        yield node
+        for child in tree.children(node):
+            yield from traversal_preorder(tree, child)
+
+
+def traversal_postorder(tree, root=None):
+    roots = tree.roots if root is None else [root]
+    for node in roots:
+        for child in tree.children(node):
+            yield from traversal_postorder(tree, child)
+        yield node
+
+
+def traversal_inorder(tree, root=None):
+    roots = tree.roots if root is None else [root]
+    for node in roots:
+        children = list(tree.children(node))
+        half = len(children) // 2
+        for child in children[:half]:
+            yield from traversal_inorder(tree, child)
+        yield node
+        for child in children[half:]:
+            yield from traversal_inorder(tree, child)
+
+
+def traversal_levelorder(tree, root=None):
+    roots = tree.roots if root is None else [root]
+    for node in roots:
+        # Sorted is a stable ordering so nodes are in tree order within
+        # the level
+        yield from sorted(list(tree.nodes(node)), key=lambda u: tree.depth(u))
+
+
+def _traversal_minlex_postorder(tree, u):
+    """
+    For a given input ID u, this function returns a tuple whose first value
+    is the minimum leaf node ID under node u, and whose second value is
+    a list containing the minlex postorder for the subtree rooted at node u.
+    The first value is needed for sorting, and the second value is what
+    finally gets returned.
+    """
+    children = tree.children(u)
+    if len(children) > 0:
+        children_return = [_traversal_minlex_postorder(tree, c) for c in children]
+        # sorts by first value, which is the minimum leaf node ID
+        children_return.sort(key=lambda x: x[0])
+        minlex_postorder = []
+        for _, child_minlex_postorder in children_return:
+            minlex_postorder.extend(child_minlex_postorder)
+        minlex_postorder.extend([u])
+        return (children_return[0][0], minlex_postorder)
+    else:
+        return (u, [u])
+
+
+def traversal_minlex_postorder(tree, root=None):
+    roots = tree.roots if root is None else [root]
+    root_lists = [_traversal_minlex_postorder(tree, node) for node in roots]
+    for _, node_list in sorted(root_lists, key=lambda x: x[0]):
+        yield from node_list
+
+
+def traversal_timeasc(tree, root=None):
+    # This is wrong, it should be
+    # nodes = sorted(tree.nodes(root), key=lambda u: (tree.time(u), u))
+    # https://github.com/tskit-dev/tskit/issues/1776
+    roots = tree.roots if root is None else [root]
+    for node in roots:
+        nodes = sorted(tree.nodes(node), key=lambda u: (tree.time(u), u))
+        yield from nodes
+
+
+def traversal_timedesc(tree, root=None):
+    # This is wrong, it should be
+    # nodes = sorted(tree.nodes(root), key=lambda u: (-tree.time(u), u))
+    # https://github.com/tskit-dev/tskit/issues/1776
+    roots = tree.roots if root is None else [root]
+    for node in roots:
+        nodes = sorted(tree.nodes(node), key=lambda u: (tree.time(u), u), reverse=True)
+        yield from nodes
+
+
+traversal_map = {
+    "preorder": traversal_preorder,
+    "postorder": traversal_postorder,
+    "inorder": traversal_inorder,
+    "levelorder": traversal_levelorder,
+    "breadthfirst": traversal_levelorder,
+    "minlex_postorder": traversal_minlex_postorder,
+    "timeasc": traversal_timeasc,
+    "timedesc": traversal_timedesc,
+}
+
+
 def insert_uniform_mutations(tables, num_mutations, nodes):
     """
     Returns n evenly mutations over the specified list of nodes.
@@ -363,6 +460,39 @@ def get_samples(ts, time=None, population=None):
     return np.array(samples)
 
 
+class TestTreeTraversals:
+    @pytest.mark.parametrize("ts", get_example_tree_sequences())
+    @pytest.mark.parametrize(
+        "order",
+        [
+            "preorder",
+            # "postorder",
+            # "inorder",
+            # "levelorder",
+            # "breadthfirst",
+            # "minlex_postorder",
+        ],
+    )
+    def test_traversals_virtual_root(self, ts, order):
+        tree = ts.first()
+        node_list2 = list(traversal_map[order](tree, tree.virtual_root))
+        node_list1 = list(tree.nodes(tree.virtual_root, order=order))
+        assert tree.virtual_root in node_list1
+        assert node_list1 == node_list2
+
+    @pytest.mark.parametrize("ts", get_example_tree_sequences())
+    @pytest.mark.parametrize("order", list(traversal_map.keys()))
+    def test_traversals(self, ts, order):
+        tree = next(ts.trees())
+        traverser = traversal_map[order]
+        node_list1 = list(tree.nodes(order=order))
+        node_list2 = list(traverser(tree))
+        assert node_list1 == node_list2
+
+    # TODO add tests for traversing from specific nodes and a couple
+    # of hand crafted examples just for sanity checking.
+
+
 class TestMRCACalculator:
     """
     Class to test the Schieber-Vishkin algorithm.
@@ -2649,75 +2779,6 @@ def verify_nx_nearest_neighbor_search(self):
         nearest_neighbor_of = [min(dist_dod[u], key=dist_dod[u].get) for u in range(3)]
         assert [2, 2, 1] == [nearest_neighbor_of[u] for u in range(3)]
 
-    def test_traversals(self):
-        for ts in get_example_tree_sequences():
-            tree = next(ts.trees())
-            self.verify_traversals(tree)
-
-            # Verify time-ordered traversals separately, because the PythonTree
-            # class does not contain time information at the moment
-            for root in tree.roots:
-                time_ordered = tree.nodes(root, order="timeasc")
-                t = tree.time(next(time_ordered))
-                for u in time_ordered:
-                    next_t = tree.time(u)
-                    assert next_t >= t
-                    t = next_t
-                time_ordered = tree.nodes(root, order="timedesc")
-                t = tree.time(next(time_ordered))
-                for u in time_ordered:
-                    next_t = tree.time(u)
-                    assert next_t <= t
-                    t = next_t
-
-    def verify_traversals(self, tree):
-        t1 = tree
-        t2 = tests.PythonTree.from_tree(t1)
-        assert list(t1.nodes()) == list(t2.nodes())
-        orders = [
-            "inorder",
-            "postorder",
-            "levelorder",
-            "breadthfirst",
-            "minlex_postorder",
-        ]
-        if tree.num_roots == 1:
-            with pytest.raises(ValueError):
-                list(t1.nodes(order="bad order"))
-            assert list(t1.nodes()) == list(t1.nodes(t1.get_root()))
-            assert list(t1.nodes()) == list(t1.nodes(t1.get_root(), "preorder"))
-            for u in t1.nodes():
-                assert list(t1.nodes(u)) == list(t2.nodes(u))
-            for test_order in orders:
-                assert sorted(list(t1.nodes())) == sorted(
-                    list(t1.nodes(order=test_order))
-                )
-                assert list(t1.nodes(order=test_order)) == list(
-                    t1.nodes(t1.get_root(), order=test_order)
-                )
-                assert list(t1.nodes(order=test_order)) == list(
-                    t1.nodes(t1.get_root(), test_order)
-                )
-                assert list(t1.nodes(order=test_order)) == list(
-                    t2.nodes(order=test_order)
-                )
-                for u in t1.nodes():
-                    assert list(t1.nodes(u, test_order)) == list(
-                        t2.nodes(u, test_order)
-                    )
-        else:
-            for test_order in orders:
-                all_nodes = []
-                for root in t1.roots:
-                    assert list(t1.nodes(root, order=test_order)) == list(
-                        t2.nodes(root, order=test_order)
-                    )
-                    all_nodes.extend(t1.nodes(root, order=test_order))
-                # minlex_postorder reorders the roots, so this last test is
-                # not appropriate
-                if test_order != "minlex_postorder":
-                    assert all_nodes == list(t1.nodes(order=test_order))
-
     def test_total_branch_length(self):
         # Note: this definition works when we have no non-sample branches.
         t1 = self.get_tree()