Use C implementation for seek_index

Closes #2696

Author: jeromekelleher

c/CHANGELOG.rst

@@ -2,6 +2,11 @@
 [1.1.2] - 2023-XX-XX
 --------------------
 
+**Performance improvements**
+
+- tsk_tree_seek is now much faster at seeking to arbitrary points along
+  the sequence from the null tree (:user:`molpopgen`, :pr:`2661`).
+
 **Features**
 
 - The struct ``tsk_treeseq_t`` now has the variables ``min_time`` and ``max_time``,
@@ -24,6 +29,8 @@
 - Add `x_table_keep_rows` methods to provide efficient in-place table subsetting
   (:user:`jeromekelleher`, :pr:`2700`).
 
+- Add `tsk_tree_seek_index` function
+
 --------------------
 [1.1.1] - 2022-07-29
 --------------------

c/tskit/trees.c

@@ -4640,7 +4640,7 @@ tsk_tree_seek_from_null(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(optio
      */
     tree_index = (tsk_id_t) tsk_search_sorted(breakpoints, num_trees + 1, x);
     if (breakpoints[tree_index] > x) {
-        tree_index -= 1;
+        tree_index--;
     }
     self->index = tree_index;
     self->interval.left = breakpoints[tree_index];

python/CHANGELOG.rst

@@ -2,6 +2,11 @@
 [0.5.5] - 2023-01-XX
 --------------------
 
+**Performance improvements**
+
+- Methods like ts.at() which seek to a specified position on the sequence from
+  a new Tree instance are now much faster (:user:`molpopgen`, :pr:`2661`).
+
 **Features**
 
 - Add ``__repr__`` for variants to return a string representation of the raw data

python/_tskitmodule.c

@@ -10658,6 +10658,29 @@ Tree_seek(Tree *self, PyObject *args)
     return ret;
 }
 
+static PyObject *
+Tree_seek_index(Tree *self, PyObject *args)
+{
+    PyObject *ret = NULL;
+    tsk_id_t index = 0;
+    int err;
+
+    if (Tree_check_state(self) != 0) {
+        goto out;
+    }
+    if (!PyArg_ParseTuple(args, "O&", tsk_id_converter, &index)) {
+        goto out;
+    }
+    err = tsk_tree_seek_index(self->tree, index, 0);
+    if (err != 0) {
+        handle_library_error(err);
+        goto out;
+    }
+    ret = Py_BuildValue("");
+out:
+    return ret;
+}
+
 static PyObject *
 Tree_clear(Tree *self)
 {
@@ -11796,6 +11819,10 @@ static PyMethodDef Tree_methods[] = {
         .ml_meth = (PyCFunction) Tree_seek,
         .ml_flags = METH_VARARGS,
         .ml_doc = "Seeks to the tree at the specified position" },
+    { .ml_name = "seek_index",
+        .ml_meth = (PyCFunction) Tree_seek_index,
+        .ml_flags = METH_VARARGS,
+        .ml_doc = "Seeks to the tree at the specified index" },
     { .ml_name = "clear",
         .ml_meth = (PyCFunction) Tree_clear,
         .ml_flags = METH_NOARGS,

python/tests/test_lowlevel.py

@@ -2968,6 +2968,16 @@ def test_seek_errors(self):
             with pytest.raises(_tskit.LibraryError):
                 tree.seek(bad_pos)
 
+    def test_seek_index_errors(self):
+        ts = self.get_example_tree_sequence()
+        tree = _tskit.Tree(ts)
+        for bad_type in ["", "x", {}]:
+            with pytest.raises(TypeError):
+                tree.seek_index(bad_type)
+        for bad_index in [-1, 10**6]:
+            with pytest.raises(_tskit.LibraryError):
+                tree.seek_index(bad_index)
+
     def test_root_threshold(self):
         for ts in self.get_example_tree_sequences():
             tree = _tskit.Tree(ts)

python/tskit/trees.py

@@ -820,7 +820,6 @@ def seek_index(self, index):
 
         .. include:: substitutions/linear_traversal_warning.rst
 
-
         :param int index: The tree index to seek to.
         :raises IndexError: If an index outside the acceptable range is provided.
         """
@@ -829,12 +828,7 @@ def seek_index(self, index):
             index += num_trees
         if index < 0 or index >= num_trees:
             raise IndexError("Index out of bounds")
-        # This should be implemented in C efficiently using the indexes.
-        # No point in complicating the current implementation by trying
-        # to seek from the correct direction.
-        self.first()
-        while self.index != index:
-            self.next()
+        self._ll_tree.seek_index(index)
 
     def seek(self, position):
         """