Refine API, add tests and fix implementation error

Author: jeromekelleher

c/tests/test_trees.c

@@ -331,7 +331,7 @@ verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents)
     uint32_t mutation_index, site_index;
     tsk_size_t k, l, tree_sites_length;
     const tsk_site_t *sites = NULL;
-    tsk_tree_t tree;
+    tsk_tree_t tree, skip_tree;
     tsk_size_t num_edges;
     tsk_size_t num_nodes = tsk_treeseq_get_num_nodes(ts);
     tsk_size_t num_sites = tsk_treeseq_get_num_sites(ts);
@@ -340,6 +340,7 @@ verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents)
 
     ret = tsk_tree_init(&tree, ts, 0);
     CU_ASSERT_EQUAL(ret, 0);
+    ret = tsk_tree_init(&skip_tree, ts, 0);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_trees(ts), num_trees);
 
     CU_ASSERT_EQUAL(tree.index, -1);
@@ -372,6 +373,22 @@ verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents)
             }
             site_index++;
         }
+        /* Check the skip tree */
+        /* printf("Call first %p\n", (void *) &skip_tree); */
+        ret = tsk_tree_first(&skip_tree);
+        /* tsk_tree_print_state(&skip_tree, stdout); */
+        CU_ASSERT_EQUAL(ret, TSK_TREE_OK);
+        /* printf("Call seek %p\n", (void *) &skip_tree); */
+        ret = tsk_tree_seek(&skip_tree, breakpoints[j], TSK_TREE_SEEK_SKIP);
+        CU_ASSERT_EQUAL(ret, 0);
+        /* tsk_tree_print_state(&skip_tree, _devnull); */
+        check_trees_equal(&tree, &skip_tree);
+        ret = tsk_tree_last(&skip_tree);
+        CU_ASSERT_EQUAL(ret, TSK_TREE_OK);
+        ret = tsk_tree_seek(&skip_tree, breakpoints[j], TSK_TREE_SEEK_SKIP);
+        CU_ASSERT_EQUAL(ret, 0);
+        /* tsk_tree_print_state(&skip_tree, _devnull); */
+        check_trees_equal(&tree, &skip_tree);
         j++;
     }
     CU_ASSERT_EQUAL(ret, 0);
@@ -381,6 +398,7 @@ verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents)
     CU_ASSERT_EQUAL(tsk_treeseq_get_sequence_length(ts), breakpoints[j]);
 
     tsk_tree_free(&tree);
+    tsk_tree_free(&skip_tree);
 
     verify_tree_pos(ts, num_trees, parents);
 }
@@ -6283,7 +6301,7 @@ test_multiroot_tree_traversal(void)
 }
 
 static void
-test_seek_multi_tree(void)
+verify_seek_multi_tree(tsk_flags_t seek_options)
 {
     int ret;
     tsk_treeseq_t ts;
@@ -6299,14 +6317,14 @@ test_seek_multi_tree(void)
     CU_ASSERT_EQUAL_FATAL(ret, 0);
 
     for (j = 0; j < num_trees; j++) {
-        ret = tsk_tree_seek(&t, breakpoints[j], 0);
+        ret = tsk_tree_seek(&t, breakpoints[j], seek_options);
         CU_ASSERT_EQUAL_FATAL(ret, 0);
         CU_ASSERT_EQUAL_FATAL(t.index, j);
         ret = tsk_tree_seek_index(&t, j, 0);
         CU_ASSERT_EQUAL_FATAL(ret, 0);
         CU_ASSERT_EQUAL_FATAL(t.index, j);
         for (k = 0; k < num_trees; k++) {
-            ret = tsk_tree_seek(&t, breakpoints[k], 0);
+            ret = tsk_tree_seek(&t, breakpoints[k], seek_options);
             CU_ASSERT_EQUAL_FATAL(ret, 0);
             CU_ASSERT_EQUAL_FATAL(t.index, k);
             ret = tsk_tree_seek_index(&t, k, 0);
@@ -6315,13 +6333,13 @@ test_seek_multi_tree(void)
         }
     }
 
-    ret = tsk_tree_seek(&t, 1.99999, 0);
+    ret = tsk_tree_seek(&t, 1.99999, seek_options);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
     CU_ASSERT_EQUAL_FATAL(t.index, 0);
-    ret = tsk_tree_seek(&t, 6.99999, 0);
+    ret = tsk_tree_seek(&t, 6.99999, seek_options);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
     CU_ASSERT_EQUAL_FATAL(t.index, 1);
-    ret = tsk_tree_seek(&t, 9.99999, 0);
+    ret = tsk_tree_seek(&t, 9.99999, seek_options);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
     CU_ASSERT_EQUAL_FATAL(t.index, 2);
 
@@ -6331,7 +6349,7 @@ test_seek_multi_tree(void)
     for (j = 0; j < num_trees; j++) {
         ret = tsk_tree_init(&t, &ts, 0);
         CU_ASSERT_EQUAL_FATAL(ret, 0);
-        ret = tsk_tree_seek(&t, breakpoints[j], 0);
+        ret = tsk_tree_seek(&t, breakpoints[j], seek_options);
         CU_ASSERT_EQUAL_FATAL(ret, 0);
         CU_ASSERT_EQUAL_FATAL(t.index, j);
         tsk_tree_free(&t);
@@ -6341,12 +6359,12 @@ test_seek_multi_tree(void)
     ret = tsk_tree_init(&t, &ts, 0);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
     for (j = 0; j < num_trees; j++) {
-        ret = tsk_tree_seek(&t, 0, 0);
+        ret = tsk_tree_seek(&t, 0, seek_options);
         CU_ASSERT_EQUAL_FATAL(ret, 0);
         ret = tsk_tree_prev(&t);
         CU_ASSERT_EQUAL_FATAL(ret, 0);
         CU_ASSERT_EQUAL_FATAL(t.index, -1);
-        ret = tsk_tree_seek(&t, breakpoints[j], 0);
+        ret = tsk_tree_seek(&t, breakpoints[j], seek_options);
         CU_ASSERT_EQUAL_FATAL(ret, 0);
         CU_ASSERT_EQUAL_FATAL(t.index, j);
     }
@@ -6355,6 +6373,13 @@ test_seek_multi_tree(void)
     tsk_treeseq_free(&ts);
 }
 
+static void
+test_seek_multi_tree(void)
+{
+    verify_seek_multi_tree(0);
+    verify_seek_multi_tree(TSK_TREE_SEEK_SKIP);
+}
+
 static void
 test_seek_errors(void)
 {

c/tskit/trees.c

@@ -6146,17 +6146,18 @@ tsk_tree_print_state(const tsk_tree_t *self, FILE *out)
     fprintf(out, "left = %f\n", self->interval.left);
     fprintf(out, "right = %f\n", self->interval.right);
     fprintf(out, "index = %lld\n", (long long) self->index);
-    fprintf(out, "node\tparent\tlchild\trchild\tlsib\trsib");
+    fprintf(out, "num_edges = %d\n", (int) self->num_edges);
+    fprintf(out, "node\tedge\tparent\tlchild\trchild\tlsib\trsib");
     if (self->options & TSK_SAMPLE_LISTS) {
         fprintf(out, "\thead\ttail");
     }
     fprintf(out, "\n");
 
     for (j = 0; j < self->num_nodes + 1; j++) {
-        fprintf(out, "%lld\t%lld\t%lld\t%lld\t%lld\t%lld", (long long) j,
-            (long long) self->parent[j], (long long) self->left_child[j],
-            (long long) self->right_child[j], (long long) self->left_sib[j],
-            (long long) self->right_sib[j]);
+        fprintf(out, "%lld\t%lld\t%lld\t%lld\t%lld\t%lld\t%lld", (long long) j,
+            (long long) self->edge[j], (long long) self->parent[j],
+            (long long) self->left_child[j], (long long) self->right_child[j],
+            (long long) self->left_sib[j], (long long) self->right_sib[j]);
         if (self->options & TSK_SAMPLE_LISTS) {
             fprintf(out, "\t%lld\t%lld\t", (long long) self->left_sample[j],
                 (long long) self->right_sample[j]);
@@ -6284,7 +6285,7 @@ tsk_tree_remove_root(tsk_tree_t *self, tsk_id_t root, tsk_id_t *restrict parent)
 }
 
 static void
-tsk_tree_remove_edge(tsk_tree_t *self, tsk_id_t p, tsk_id_t c)
+tsk_tree_remove_edge(tsk_tree_t *self, tsk_id_t p, tsk_id_t c, tsk_id_t TSK_UNUSED(edge_id))
 {
     tsk_id_t *restrict parent = self->parent;
     tsk_size_t *restrict num_samples = self->num_samples;
@@ -6295,6 +6296,7 @@ tsk_tree_remove_edge(tsk_tree_t *self, tsk_id_t p, tsk_id_t c)
     tsk_id_t path_end = TSK_NULL;
     bool path_end_was_root = false;
 
+    /* printf("%p remove edge (%d) %d->%d\n", (void *) self, edge_id, (int) c, (int) p); */
 #define POTENTIAL_ROOT(U) (num_samples[U] >= root_threshold)
 
     tsk_tree_remove_branch(self, p, c, parent);
@@ -6336,6 +6338,8 @@ tsk_tree_insert_edge(tsk_tree_t *self, tsk_id_t p, tsk_id_t c, tsk_id_t edge_id)
     tsk_id_t path_end = TSK_NULL;
     bool path_end_was_root = false;
 
+    /* printf("%p insert edge (%d) %d->%d\n", (void *) self, edge_id, (int) c, (int) p); */
+
 #define POTENTIAL_ROOT(U) (num_samples[U] >= root_threshold)
 
     if (!(self->options & TSK_NO_SAMPLE_COUNTS)) {
@@ -6425,7 +6429,7 @@ tsk_tree_next(tsk_tree_t *self)
     if (valid) {
         for (j = tree_pos.out.start; j != tree_pos.out.stop; j++) {
             e = tree_pos.out.order[j];
-            tsk_tree_remove_edge(self, edge_parent[e], edge_child[e]);
+            tsk_tree_remove_edge(self, edge_parent[e], edge_child[e], e);
         }
 
         for (j = tree_pos.in.start; j != tree_pos.in.stop; j++) {
@@ -6457,7 +6461,7 @@ tsk_tree_prev(tsk_tree_t *self)
     if (valid) {
         for (j = tree_pos.out.start; j != tree_pos.out.stop; j--) {
             e = tree_pos.out.order[j];
-            tsk_tree_remove_edge(self, edge_parent[e], edge_child[e]);
+            tsk_tree_remove_edge(self, edge_parent[e], edge_child[e], e);
         }
 
         for (j = tree_pos.in.start; j != tree_pos.in.stop; j--) {
@@ -6556,9 +6560,9 @@ tsk_tree_seek_forward(tsk_tree_t *self, tsk_id_t index)
     for (j = tree_pos.out.start; j != tree_pos.out.stop; j++) {
         e = tree_pos.out.order[j];
         e_left = edge_left[e];
-        if (e_left <= old_right) {
+        if (e_left < old_right) {
             tsk_bug_assert(edge_parent[e] != TSK_NULL);
-            tsk_tree_remove_edge(self, edge_parent[e], edge_child[e]);
+            tsk_tree_remove_edge(self, edge_parent[e], edge_child[e], e);
         }
         tsk_bug_assert(e_left < interval_left);
     }
@@ -6600,7 +6604,7 @@ tsk_tree_seek_backward(tsk_tree_t *self, tsk_id_t index)
         e_right = edge_right[e];
         if (e_right >= old_right) {
             tsk_bug_assert(edge_parent[e] != TSK_NULL);
-            tsk_tree_remove_edge(self, edge_parent[e], edge_child[e]);
+            tsk_tree_remove_edge(self, edge_parent[e], edge_child[e], e);
         }
         tsk_bug_assert(e_right > interval_right);
     }
@@ -6709,7 +6713,7 @@ tsk_tree_seek(tsk_tree_t *self, double x, tsk_flags_t options)
     if (self->index == -1) {
         ret = tsk_tree_seek_from_null(self, x, options);
     } else {
-        if (options & TSK_TREE_SEEK_ENABLE_SKIPPING) {
+        if (options & TSK_TREE_SEEK_SKIP) {
             ret = tsk_tree_seek_skip(self, x);
         } else {
             ret = tsk_tree_seek_linear(self, x);

c/tskit/trees.h

@@ -1275,7 +1275,7 @@ int tsk_tree_copy(const tsk_tree_t *self, tsk_tree_t *dest, tsk_flags_t options)
 
     @ingroup TREE_API_SEEKING_GROUP
 */
-#define TSK_TREE_SEEK_ENABLE_SKIPPING (1 << 0)
+#define TSK_TREE_SEEK_SKIP (1 << 0)
 
 /**
 @brief Seek to the first tree in the sequence.
@@ -1286,7 +1286,7 @@ tree sequence.
 @endrst
 
 @param self A pointer to an initialised tsk_tree_t object.
-@return Return on success; or a negative value if an error occurs.
+@return Return TSK_TREE_OK on success; or a negative value if an error occurs.
 */
 int tsk_tree_first(tsk_tree_t *self);
 
@@ -1384,13 +1384,13 @@ Seeking to a position currently covered by the tree is
 a constant time operation.
 
 Seeking to a position from a non-null tree use a linear time
-algorithm by default, unless the option :c:macro:`TSK_TREE_SEEK_ENABLE_SKIPPING`
+algorithm by default, unless the option :c:macro:`TSK_TREE_SKIP`
 is specified. In this case, a faster algorithm is employed which skips
 to the target tree by removing and adding the minimal number of edges
 possible. However, this approach does not guarantee that edges are
 inserted and removed in time-sorted order.
 
-.. warning:: Using the :c:macro:`TSK_TREE_SEEK_ENABLE_SKIPPING` option
+.. warning:: Using the :c:macro:`TSK_TREE_SEEK_SKIP` option
     may lead to edges not being inserted or removed in time-sorted order.
 
 @endrst

python/_tskitmodule.c

@@ -12104,20 +12104,19 @@ Tree_seek(Tree *self, PyObject *args)
 {
     PyObject *ret = NULL;
     tsk_flags_t options = 0;
-    int enable_skipping = true;
+    int skip = false;
     double position;
     int err;
 
-    if (enable_skipping) {
-        options |= TSK_TREE_SEEK_ENABLE_SKIPPING;
-    }
-
     if (Tree_check_state(self) != 0) {
         goto out;
     }
-    if (!PyArg_ParseTuple(args, "d", &position)) {
+    if (!PyArg_ParseTuple(args, "d|i", &position, &skip)) {
         goto out;
     }
+    if (skip) {
+        options |= TSK_TREE_SEEK_SKIP;
+    }
     err = tsk_tree_seek(self->tree, position, options);
     if (err != 0) {
         handle_library_error(err);
@@ -12133,15 +12132,20 @@ Tree_seek_index(Tree *self, PyObject *args)
 {
     PyObject *ret = NULL;
     tsk_id_t index = 0;
+    tsk_flags_t options = 0;
+    int skip = false;
     int err;
 
     if (Tree_check_state(self) != 0) {
         goto out;
     }
-    if (!PyArg_ParseTuple(args, "O&", tsk_id_converter, &index)) {
+    if (!PyArg_ParseTuple(args, "O&|i", tsk_id_converter, &index, &skip)) {
         goto out;
     }
-    err = tsk_tree_seek_index(self->tree, index, 0);
+    if (skip) {
+        options |= TSK_TREE_SEEK_SKIP;
+    }
+    err = tsk_tree_seek_index(self->tree, index, options);
     if (err != 0) {
         handle_library_error(err);
         goto out;

python/tests/test_highlevel.py

@@ -4373,11 +4373,11 @@ def test_nonbinary(self):
 def assert_trees_identical(t1, t2):
     assert t1.tree_sequence == t2.tree_sequence
     assert t1.index == t2.index
-    assert np.all(t1.parent_array == t2.parent_array)
-    assert np.all(t1.left_child_array == t2.left_child_array)
-    assert np.all(t1.left_sib_array == t2.left_sib_array)
-    assert np.all(t1.right_child_array == t2.right_child_array)
-    assert np.all(t1.right_sib_array == t2.right_sib_array)
+    assert_array_equal(t1.parent_array, t2.parent_array)
+    assert_array_equal(t1.left_child_array, t2.left_child_array)
+    assert_array_equal(t1.left_sib_array, t2.left_sib_array)
+    assert_array_equal(t1.right_child_array, t2.right_child_array)
+    assert_array_equal(t1.right_sib_array, t2.right_sib_array)
 
 
 def assert_same_tree_different_order(t1, t2):
@@ -4431,8 +4431,9 @@ def get_tree_pair(self):
         ts = self.ts()
         t1 = tskit.Tree(ts)
         t2 = tskit.Tree(ts)
-        # Note: for development we can monkeypatch in the Python implementation
-        # above like this:
+        # # Note: for development we can monkeypatch in the Python implementation
+        # # above like this:
+        # import functools
         # t2.seek = functools.partial(seek, t2)
         return t1, t2
 
@@ -4447,34 +4448,36 @@ def test_index_from_different_directions(self, index):
             t2.prev()
         assert_same_tree_different_order(t1, t2)
 
-    @pytest.mark.skip()
-    #   @pytest.mark.parametrize("position", [0, 1, 2, 3])
-    def test_seek_from_null(self, position):
+    @pytest.mark.parametrize("position", [0, 1, 2, 3])
+    @pytest.mark.parametrize("skip", [False, True])
+    def test_seek_from_null(self, position, skip):
         t1, t2 = self.get_tree_pair()
         t1.clear()
         t1.seek(position)
         t2.first()
-        t2.seek(position)
+        t2.seek(position, skip=skip)
         assert_trees_identical(t1, t2)
 
     @pytest.mark.parametrize("index", range(3))
-    def test_seek_next_tree(self, index):
+    @pytest.mark.parametrize("skip", [False, True])
+    def test_seek_next_tree(self, index, skip):
         t1, t2 = self.get_tree_pair()
         while t1.index != index:
             t1.next()
             t2.next()
         t1.next()
-        t2.seek(index + 1)
+        t2.seek(index + 1, skip=skip)
         assert_trees_identical(t1, t2)
 
     @pytest.mark.parametrize("index", [3, 2, 1])
-    def test_seek_prev_tree(self, index):
+    @pytest.mark.parametrize("skip", [False, True])
+    def test_seek_prev_tree(self, index, skip):
         t1, t2 = self.get_tree_pair()
         while t1.index != index:
             t1.prev()
             t2.prev()
         t1.prev()
-        t2.seek(index - 1)
+        t2.seek(index - 1, skip=skip)
         assert_trees_identical(t1, t2)
 
     def test_seek_1_from_0(self):
@@ -4515,15 +4518,13 @@ def test_seek_3_from_null_prev(self):
         t2.prev()
         assert_trees_identical(t1, t2)
 
-    @pytest.mark.skip()
     def test_seek_3_from_0(self):
         t1, t2 = self.get_tree_pair()
         t1.last()
         t2.first()
         t2.seek(3)
         assert_trees_identical(t1, t2)
 
-    @pytest.mark.skip()
     def test_seek_0_from_3(self):
         t1, t2 = self.get_tree_pair()
         t1.last()