Add link-ancestors to TreeSequence and ImmutableTableCollection

python/_tskitmodule.c

@@ -5344,6 +5344,69 @@ TreeSequence_dump_tables(TreeSequence *self, PyObject *args, PyObject *kwds)
     return ret;
 }
 
+static PyObject *
+TreeSequence_link_ancestors(TreeSequence *self, PyObject *args, PyObject *kwds)
+{
+    int err;
+    PyObject *ret = NULL;
+    PyObject *samples = NULL;
+    PyObject *ancestors = NULL;
+    PyArrayObject *samples_array = NULL;
+    PyArrayObject *ancestors_array = NULL;
+    npy_intp *shape;
+    tsk_size_t num_samples, num_ancestors;
+    EdgeTable *result = NULL;
+    PyObject *result_args = NULL;
+    static char *kwlist[] = { "samples", "ancestors", NULL };
+
+    if (TreeSequence_check_state(self) != 0) {
+        goto out;
+    }
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO", kwlist, &samples, &ancestors)) {
+        goto out;
+    }
+
+    samples_array = (PyArrayObject *) PyArray_FROMANY(
+        samples, NPY_INT32, 1, 1, NPY_ARRAY_IN_ARRAY);
+    if (samples_array == NULL) {
+        goto out;
+    }
+    shape = PyArray_DIMS(samples_array);
+    num_samples = (tsk_size_t) shape[0];
+
+    ancestors_array = (PyArrayObject *) PyArray_FROMANY(
+        ancestors, NPY_INT32, 1, 1, NPY_ARRAY_IN_ARRAY);
+    if (ancestors_array == NULL) {
+        goto out;
+    }
+    shape = PyArray_DIMS(ancestors_array);
+    num_ancestors = (tsk_size_t) shape[0];
+
+    result_args = PyTuple_New(0);
+    if (result_args == NULL) {
+        goto out;
+    }
+    result = (EdgeTable *) PyObject_CallObject((PyObject *) &EdgeTableType, result_args);
+    if (result == NULL) {
+        goto out;
+    }
+    err = tsk_table_collection_link_ancestors(self->tree_sequence->tables,
+        PyArray_DATA(samples_array), num_samples, PyArray_DATA(ancestors_array),
+        num_ancestors, 0, result->table);
+    if (err != 0) {
+        handle_library_error(err);
+        goto out;
+    }
+    ret = (PyObject *) result;
+    result = NULL;
+out:
+    Py_XDECREF(samples_array);
+    Py_XDECREF(ancestors_array);
+    Py_XDECREF(result);
+    Py_XDECREF(result_args);
+    return ret;
+}
+
 static PyObject *
 TreeSequence_load(TreeSequence *self, PyObject *args, PyObject *kwds)
 {
@@ -8528,6 +8591,10 @@ static PyMethodDef TreeSequence_methods[] = {
         .ml_meth = (PyCFunction) TreeSequence_dump_tables,
         .ml_flags = METH_VARARGS | METH_KEYWORDS,
         .ml_doc = "Dumps the tree sequence to the specified set of tables" },
+    { .ml_name = "link_ancestors",
+        .ml_meth = (PyCFunction) TreeSequence_link_ancestors,
+        .ml_flags = METH_VARARGS | METH_KEYWORDS,
+        .ml_doc = "Returns an EdgeTable linking the specified samples and ancestors." },
     { .ml_name = "get_node",
         .ml_meth = (PyCFunction) TreeSequence_get_node,
         .ml_flags = METH_VARARGS,

python/tests/test_topology.py

@@ -4901,6 +4901,11 @@ def do_map(self, ts, ancestors, samples=None, compare_lib=True):
         if compare_lib:
             lib_result = ts.dump_tables().link_ancestors(samples, ancestors)
             assert ancestor_table == lib_result
+            ts_result = ts.link_ancestors(samples, ancestors)
+            assert ancestor_table == ts_result
+            if _tskit.HAS_NUMPY_2:
+                tables_result = ts.tables.link_ancestors(samples, ancestors)
+                assert ancestor_table == tables_result
         return ancestor_table
 
     def test_deprecated_name(self):
@@ -4914,6 +4919,11 @@ def test_deprecated_name(self):
         tss = s.link_ancestors()
         lib_result = ts.dump_tables().map_ancestors(samples, ancestors)
         assert tss == lib_result
+        ts_result = ts.link_ancestors(samples, ancestors)
+        assert tss == ts_result
+        if _tskit.HAS_NUMPY_2:
+            immutable_result = ts.tables.map_ancestors(samples, ancestors)
+            assert tss == immutable_result
         assert list(tss.parent) == [8, 8, 8, 8, 8]
         assert list(tss.child) == [0, 1, 2, 3, 4]
         assert all(tss.left) == 0

python/tskit/tables.py

@@ -4780,6 +4780,21 @@ def __str__(self):
             ]
         )
 
+    def link_ancestors(self, samples, ancestors):
+        """
+        See :meth:`TableCollection.link_ancestors`.
+        """
+        samples = util.safe_np_int_cast(samples, np.int32)
+        ancestors = util.safe_np_int_cast(ancestors, np.int32)
+        ll_edge_table = self._llts.link_ancestors(samples, ancestors)
+        return EdgeTable(ll_table=ll_edge_table)
+
+    def map_ancestors(self, *args, **kwargs):
+        """
+        Deprecated alias for :meth:`link_ancestors`.
+        """
+        return self.link_ancestors(*args, **kwargs)
+
     _MUTATOR_METHODS = {
         "clear",
         "sort",
@@ -4803,8 +4818,6 @@ def __str__(self):
         "ibd_segments",
         "fromdict",
         "simplify",
-        "link_ancestors",
-        "map_ancestors",
     }
 
     def copy(self):

python/tskit/trees.py

@@ -4372,6 +4372,22 @@ def dump_tables(self):
         self._ll_tree_sequence.dump_tables(ll_tables)
         return tables.TableCollection(ll_tables=ll_tables)
 
+    def link_ancestors(self, samples, ancestors):
+        """
+        Equivalent to :meth:`TableCollection.link_ancestors`; see that method for full
+        documentation and parameter semantics.
+
+        :param list[int] samples: Node IDs to retain as samples.
+        :param list[int] ancestors: Node IDs to treat as ancestors.
+        :return: An :class:`tables.EdgeTable` containing the genealogical links between
+            the supplied ``samples`` and ``ancestors``.
+        :rtype: tables.EdgeTable
+        """
+        samples = util.safe_np_int_cast(samples, np.int32)
+        ancestors = util.safe_np_int_cast(ancestors, np.int32)
+        ll_edge_table = self._ll_tree_sequence.link_ancestors(samples, ancestors)
+        return tables.EdgeTable(ll_table=ll_edge_table)
+
     def dump_text(
         self,
         nodes=None,