Cache references to edges_parent & edges_child

And remove the need to specify which nodes to constrain, which was never used anyway.

Author: hyanwong

requirements.txt

@@ -1,4 +1,4 @@
-tskit>=0.4.0
+tskit>=0.5.2
 tsinfer>=0.2.0
 flake8
 numpy

tests/test_functions.py

@@ -1530,18 +1530,7 @@ def test_constrain_ages_topo(self):
         ts = utility_functions.two_tree_ts()
         post_mn = np.array([0.0, 0.0, 0.0, 2.0, 1.0, 3.0])
         eps = 1e-6
-        nodes_to_date = np.array([3, 4, 5])
-        constrained_ages = constrain_ages_topo(ts, post_mn, eps, nodes_to_date)
-        assert np.array_equal(
-            np.array([0.0, 0.0, 0.0, 2.0, 2.000001, 3.0]), constrained_ages
-        )
-
-    def test_constrain_ages_topo_no_nodes_to_date(self):
-        ts = utility_functions.two_tree_ts()
-        post_mn = np.array([0.0, 0.0, 0.0, 2.0, 1.0, 3.0])
-        eps = 1e-6
-        nodes_to_date = None
-        constrained_ages = constrain_ages_topo(ts, post_mn, eps, nodes_to_date)
+        constrained_ages = constrain_ages_topo(ts, post_mn, eps)
         assert np.array_equal(
             np.array([0.0, 0.0, 0.0, 2.0, 2.000001, 3.0]), constrained_ages
         )

tsdate/core.py

@@ -936,22 +936,23 @@ def constrain_ages_topo(ts, node_times, eps, progress=False):
     If node_times violate topology, return increased node_times so that each node is
     guaranteed to be older than any of its their children.
     """
-    tables = ts.tables
+    edges_parent = ts.edges_parent
+    edges_child = ts.edges_child
+
     new_node_times = np.copy(node_times)
     # Traverse through the ARG, ensuring children come before parents.
     # This can be done by iterating over groups of edges with the same parent
-    new_parent_edge_idx = np.concatenate(
-        (
-            [0],
-            np.where(np.diff(tables.edges.parent) != 0)[0] + 1,
-            [tables.edges.num_rows],
-        )
-    )
-    for edges_start, edges_end in zip(
-        new_parent_edge_idx[:-1], new_parent_edge_idx[1:]
+    new_parent_edge_idx = np.where(np.diff(edges_parent) != 0)[0] + 1
+    for edges_start, edges_end in tqdm(
+        zip(
+            itertools.chain([0], new_parent_edge_idx),
+            itertools.chain(new_parent_edge_idx, [len(edges_parent)]),
+        ),
+        desc="Constrain Ages",
+        disable=not progress,
     ):
-        parent = tables.edges.parent[edges_start]
-        child_ids = tables.edges.child[edges_start:edges_end]  # May contain dups
+        parent = edges_parent[edges_start]
+        child_ids = edges_child[edges_start:edges_end]  # May contain dups
         oldest_child_time = np.max(new_node_times[child_ids])
         if oldest_child_time >= new_node_times[parent]:
             new_node_times[parent] = oldest_child_time + eps