From 74191c0a071df3f8646fc44769ebc03e7c5fbf82 Mon Sep 17 00:00:00 2001 From: Ben Jeffery Date: Tue, 4 Nov 2025 01:20:26 +0000 Subject: [PATCH] Add link-ancestors to TreeSequence and ImmutableTableCollection --- python/_tskitmodule.c | 67 +++++++++++++++++++++++++++++++++++ python/tests/test_topology.py | 10 ++++++ python/tskit/tables.py | 17 +++++++-- python/tskit/trees.py | 16 +++++++++ 4 files changed, 108 insertions(+), 2 deletions(-) diff --git a/python/_tskitmodule.c b/python/_tskitmodule.c index 78cb9f7c8e..3144b1bbfb 100644 --- a/python/_tskitmodule.c +++ b/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, diff --git a/python/tests/test_topology.py b/python/tests/test_topology.py index b693743a2a..28c1d95c93 100644 --- a/python/tests/test_topology.py +++ b/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 diff --git a/python/tskit/tables.py b/python/tskit/tables.py index cab3407c27..d61f864593 100644 --- a/python/tskit/tables.py +++ b/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): diff --git a/python/tskit/trees.py b/python/tskit/trees.py index 1e31048075..6697665ded 100644 --- a/python/tskit/trees.py +++ b/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,