Add link-ancestors to TreeSequence and ImmutableTableCollection (#3312)

Author: benjeffery

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)
 {
@@ -8624,6 +8687,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_highlevel.py

@@ -163,6 +163,47 @@ def simplify_tree_sequence(ts, samples, filter_sites=True):
     return s.simplify()
 
 
+@pytest.mark.parametrize("ts", tsutil.get_example_tree_sequences())
+class TestLinkAncestorsExamples:
+    def test_link_ancestors_runs_and_is_sane(self, ts):
+        # Can't link ancestors when edges have metadata.
+        if ts.tables.edges.metadata_schema != tskit.MetadataSchema(schema=None):
+            pytest.skip("link_ancestors does not support edges with metadata")
+
+        samples = ts.samples()
+        if len(samples) == 0:
+            pytest.skip("Tree sequence has no samples")
+
+        # Prefer internal nodes as ancestors; fall back to samples if none.
+        ancestor_nodes = [u.id for u in ts.nodes() if not u.is_sample()]
+        if len(ancestor_nodes) == 0:
+            ancestor_nodes = list(samples)
+
+        # Keep argument sizes modest for large examples.
+        samples = samples[: min(len(samples), 10)]
+        ancestors = ancestor_nodes[: min(len(ancestor_nodes), 10)]
+
+        result = ts.link_ancestors(samples, ancestors)
+        assert isinstance(result, tskit.EdgeTable)
+
+        # Basic invariants on the returned table.
+        assert np.all(result.left >= 0)
+        assert np.all(result.right <= ts.sequence_length)
+        if result.num_rows > 0:
+            assert np.all(result.left < result.right)
+            assert set(result.parent).issubset(set(range(ts.num_nodes)))
+            assert set(result.child).issubset(set(range(ts.num_nodes)))
+
+        # Parity with mutable TableCollection implementation.
+        mutable_result = ts.dump_tables().link_ancestors(samples, ancestors)
+        assert result == mutable_result
+
+        # Parity with immutable TableCollection, when available.
+        if getattr(_tskit, "HAS_NUMPY_2", False):
+            immutable_result = ts.tables.link_ancestors(samples, ancestors)
+            assert result == immutable_result
+
+
 def oriented_forests(n):
     """
     Implementation of Algorithm O from TAOCP section 7.2.1.6.

python/tests/test_immutable_table_collection.py

@@ -128,6 +128,27 @@ def test_str_contains_identifier(self, ts):
         s = str(immutable)
         assert "ImmutableTableCollection" in s
 
+    def test_link_ancestors_parity(self, ts):
+        # Can't link ancestors when edges have metadata.
+        if ts.tables.edges.metadata_schema != tskit.MetadataSchema(schema=None):
+            pytest.skip("link_ancestors does not support edges with metadata")
+
+        mutable, immutable = get_mutable_and_immutable(ts)
+        samples = ts.samples()
+        if len(samples) == 0:
+            pytest.skip("Tree sequence has no samples")
+
+        ancestor_nodes = [u.id for u in ts.nodes() if not u.is_sample()]
+        if len(ancestor_nodes) == 0:
+            ancestor_nodes = list(samples)
+
+        samples = samples[: min(len(samples), 10)]
+        ancestors = ancestor_nodes[: min(len(ancestor_nodes), 10)]
+
+        mutable_result = mutable.link_ancestors(samples, ancestors)
+        immutable_result = immutable.link_ancestors(samples, ancestors)
+        assert mutable_result == immutable_result
+
 
 @pytest.mark.parametrize("ts", tsutil.get_example_tree_sequences())
 class TestTablesParity:

python/tests/test_lowlevel.py

@@ -1495,6 +1495,47 @@ def test_mean_descendants(self):
             A = ts.mean_descendants([focal[2:], focal[:2]])
             assert A.shape == (ts.get_num_nodes(), 2)
 
+    def test_link_ancestors_bad_args(self):
+        ts = self.get_example_tree_sequence()
+        with pytest.raises(TypeError):
+            ts.link_ancestors()
+        with pytest.raises(TypeError):
+            ts.link_ancestors([0, 1])
+        with pytest.raises(ValueError):
+            ts.link_ancestors(samples=[0, 1], ancestors="sdf")
+        with pytest.raises(ValueError):
+            ts.link_ancestors(samples="sdf", ancestors=[0, 1])
+        with pytest.raises(_tskit.LibraryError):
+            ts.link_ancestors(samples=[0, 1], ancestors=[ts.get_num_nodes(), -1])
+        with pytest.raises(_tskit.LibraryError):
+            ts.link_ancestors(samples=[0, -1], ancestors=[0])
+
+    def test_link_ancestors(self):
+        # Check that the low-level method runs and does not mutate the tree sequence
+        # and that it matches the TableCollection implementation.
+        high_ts = msprime.simulate(4, random_seed=1)
+        ts = high_ts.ll_tree_sequence
+        samples = list(range(ts.get_num_samples()))
+        ancestors = list(range(ts.get_num_nodes()))
+        num_edges_before = ts.get_num_edges()
+        edges = ts.link_ancestors(samples, ancestors)
+        assert isinstance(edges, _tskit.EdgeTable)
+        assert edges.num_rows >= 0
+        if edges.num_rows > 0:
+            assert np.all(edges.left >= 0)
+            assert np.all(edges.right <= ts.get_sequence_length())
+            assert np.all(edges.left < edges.right)
+            assert np.all(edges.parent >= 0)
+            assert np.all(edges.parent < ts.get_num_nodes())
+            assert np.all(edges.child >= 0)
+            assert np.all(edges.child < ts.get_num_nodes())
+        assert ts.get_num_edges() == num_edges_before
+
+        # Parity with low-level TableCollection.link_ancestors
+        tc = high_ts.dump_tables()._ll_tables
+        edges_from_tables = tc.link_ancestors(samples, ancestors)
+        assert edges.equals(edges_from_tables)
+
     def test_metadata_schemas(self):
         tables = _tskit.TableCollection(1.0)
         # Set the schema

python/tests/test_topology.py

@@ -4901,6 +4901,10 @@ 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
+            tables_result = ts.tables.link_ancestors(samples, ancestors)
+            assert ancestor_table == tables_result
         return ancestor_table
 
     def test_deprecated_name(self):
@@ -4914,6 +4918,10 @@ 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
+        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

@@ -4393,6 +4393,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,