Numpy versions of timeasc and timedesc

Author: jeromekelleher

c/tests/test_trees.c

@@ -1,7 +1,7 @@
 /*
  * MIT License
  *
- * Copyright (c) 2019 Tskit Developers
+ * Copyright (c) 2019-2021 Tskit Developers
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal

python/tests/test_topology.py

@@ -6459,8 +6459,9 @@ def verify(self, ts):
             assert tree.is_internal(tree.virtual_root)
             assert not tree.is_leaf(tree.virtual_root)
             assert not tree.is_sample(tree.virtual_root)
-            # The mrca of the virtual_root and anything except itself is -1
-            assert tree.mrca(0, tree.virtual_root) == tskit.NULL
+            # The mrca of the virtual_root and anything is itself
+            assert tree.mrca(0, tree.virtual_root) == tree.virtual_root
+            assert tree.mrca(tree.virtual_root, 0) == tree.virtual_root
             assert tree.mrca(tree.virtual_root, tree.virtual_root) == tree.virtual_root
             # The virtual_root is a descendant of nothing other than itself
             assert not tree.is_descendant(0, tree.virtual_root)

python/tskit/trees.py

@@ -2097,6 +2097,22 @@ def preorder(self, u=NULL):
     def postorder(self, u=NULL):
         return self._ll_tree.get_postorder(u)
 
+    def timeasc(self, u=NULL):
+        nodes = self.preorder(u)
+        is_virtual_root = u == self.virtual_root
+        time = self.tree_sequence.tables.nodes.time
+        if is_virtual_root:
+            # We could avoid creating this array if we wanted to, but
+            # it's not that often people will be using this with the
+            # virtual_root as an argument, so doesn't seem worth
+            # the complexity
+            time = np.append(time, [np.inf])
+        order = np.lexsort([nodes, time[nodes]])
+        return nodes[order]
+
+    def timedesc(self, u=NULL):
+        return self.timeasc(u)[::-1]
+
     def _preorder_traversal(self, root):
         return map(int, self.preorder(root))
 
@@ -2137,16 +2153,13 @@ def _timeasc_traversal(self, root):
         """
         Sorts by increasing time but falls back to increasing ID for equal times.
         """
-        # TODO implement with numpy?
-        yield from sorted(self.nodes(root), key=lambda u: (self.time(u), u))
+        yield from self.timeasc(root)
 
     def _timedesc_traversal(self, root):
         """
         The reverse of timeasc.
         """
-        yield from sorted(
-            self.nodes(root), key=lambda u: (self.time(u), u), reverse=True
-        )
+        yield from self.timedesc(root)
 
     def _minlex_postorder_traversal(self, root):
         """