Add individuals nodes method

Author: benjeffery

python/CHANGELOG.rst

@@ -1,3 +1,13 @@
+--------------------
+[0.6.4] - 2025-XX-XX
+--------------------
+
+**Features**
+
+- Add ``TreeSequence.individuals_nodes`` attribute to return the nodes
+  associated with each individual as a numpy array.
+  (:user:`benjeffery`, :pr:`3153`)
+
 --------------------
 [0.6.3] - 2025-04-28
 --------------------

python/_tskitmodule.c

@@ -8725,6 +8725,61 @@ TreeSequence_get_individuals_time(TreeSequence *self)
     return ret;
 }
 
+static PyObject *
+TreeSequence_get_individuals_nodes(TreeSequence *self)
+{
+    PyObject *ret = NULL;
+    PyArrayObject *ret_array = NULL;
+    npy_intp dims[2];
+    tsk_size_t ploidy;
+    tsk_size_t max_ploidy = 0;
+    tsk_id_t *node_mem;
+    tsk_size_t j;
+    tsk_size_t num_individuals;
+    tsk_id_t *const *individual_nodes;
+    const tsk_size_t *individual_nodes_length;
+
+    if (TreeSequence_check_state(self) != 0) {
+        goto out;
+    }
+
+    num_individuals = tsk_treeseq_get_num_individuals(self->tree_sequence);
+    individual_nodes = self->tree_sequence->individual_nodes;
+    individual_nodes_length = self->tree_sequence->individual_nodes_length;
+
+    for (tsk_id_t i = 0; i < (tsk_id_t) num_individuals; i++) {
+        ploidy = individual_nodes_length[i];
+        if (ploidy > max_ploidy) {
+            max_ploidy = ploidy;
+        }
+    }
+
+    dims[0] = (npy_intp) num_individuals;
+    dims[1] = (npy_intp) max_ploidy;
+    ret_array = (PyArrayObject *) PyArray_SimpleNew(2, dims, NPY_INT32);
+    if (ret_array == NULL) {
+        goto out;
+    }
+
+    /* Fill with -1 (TSK_NULL) */
+    node_mem = (tsk_id_t *) PyArray_DATA(ret_array);
+    memset(node_mem, 0xFF, PyArray_NBYTES(ret_array));
+
+    for (tsk_id_t i = 0; i < (tsk_id_t) num_individuals; i++) {
+        ploidy = individual_nodes_length[i];
+        for (j = 0; j < ploidy; j++) {
+            node_mem[i * max_ploidy + j] = individual_nodes[i][j];
+        }
+    }
+
+    ret = (PyObject *) ret_array;
+    ret_array = NULL;
+
+out:
+    Py_XDECREF(ret_array);
+    return ret;
+}
+
 static PyObject *
 TreeSequence_genealogical_nearest_neighbours(
     TreeSequence *self, PyObject *args, PyObject *kwds)
@@ -11425,6 +11480,10 @@ static PyMethodDef TreeSequence_methods[] = {
         .ml_meth = (PyCFunction) TreeSequence_get_individuals_time,
         .ml_flags = METH_NOARGS,
         .ml_doc = "Returns the vector of per-individual times." },
+    { .ml_name = "get_individuals_nodes",
+        .ml_meth = (PyCFunction) TreeSequence_get_individuals_nodes,
+        .ml_flags = METH_NOARGS,
+        .ml_doc = "Returns an array of the node ids for each individual" },
     { .ml_name = "genealogical_nearest_neighbours",
         .ml_meth = (PyCFunction) TreeSequence_genealogical_nearest_neighbours,
         .ml_flags = METH_VARARGS | METH_KEYWORDS,

python/tests/test_highlevel.py

@@ -5491,3 +5491,96 @@ def test_error_if_no_schema(self, table_name):
         ts = msprime.simulate(10)
         with pytest.raises(NotImplementedError):
             getattr(ts, f"{table_name}_metadata")
+
+
+class TestIndividualsNodes:
+    def test_basic_individuals_nodes(self, tmp_path):
+        # Create a basic tree sequence with two individuals
+        tables = tskit.TableCollection(sequence_length=100)
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=1)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=1)
+        ts = tables.tree_sequence()
+
+        result = ts.individuals_nodes
+        assert result.shape == (2, 2)
+        assert_array_equal(result, [[0, 1], [2, 3]])
+
+    def test_variable_ploidy(self, tmp_path):
+        tables = tskit.TableCollection(sequence_length=100)
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")  # Diploid
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")  # Haploid
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")  # Triploid
+
+        # Diploid individual
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
+
+        # Haploid individual
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=1)
+
+        # Triploid individual
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=2)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=2)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=2)
+
+        ts = tables.tree_sequence()
+
+        result = ts.individuals_nodes
+
+        assert result.shape == (3, 3)
+
+        expected = np.array(
+            [[0, 1, -1], [2, -1, -1], [3, 4, 5]]  # Diploid  # Haploid  # Triploid
+        )
+        assert_array_equal(result, expected)
+
+    def test_no_individuals(self):
+        tables = tskit.TableCollection(sequence_length=100)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        ts = tables.tree_sequence()
+
+        result = ts.individuals_nodes
+        expected = np.array([], dtype=np.int32).reshape(0, 0)
+        assert result.shape == (0, 0)
+        assert_array_equal(result, expected)
+
+    def test_no_nodes_with_individuals(self):
+        tables = tskit.TableCollection(sequence_length=100)
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")
+        # Node without individual reference
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        ts = tables.tree_sequence()
+
+        result = ts.individuals_nodes
+        expected = np.array([[]])
+        assert result.shape == (1, 0)
+        assert_array_equal(result, expected)
+
+    def test_individual_with_no_nodes(self):
+        tables = tskit.TableCollection(sequence_length=100)
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")
+        # Only add nodes for first individual
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
+        ts = tables.tree_sequence()
+
+        result = ts.individuals_nodes
+        expected = np.array([[0], [-1]])
+        assert result.shape == (2, 1)
+        assert_array_equal(result, expected)
+
+    def test_mixed_sample_status(self):
+        tables = tskit.TableCollection(sequence_length=100)
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
+        tables.nodes.add_row(flags=0, time=0, individual=0)
+        ts = tables.tree_sequence()
+
+        result = ts.individuals_nodes
+        expected = np.array([[0, 1]])
+        assert result.shape == (1, 2)
+        assert_array_equal(result, expected)

python/tests/test_lowlevel.py

@@ -1949,6 +1949,31 @@ def test_array_lifetime(self, name, ts_fixture):
         a2[:] = 0
         assert a3 is not a2
 
+    def test_individuals_nodes(self, ts_fixture):
+        ts_fixture = ts_fixture.ll_tree_sequence
+
+        # Properties
+        a = ts_fixture.get_individuals_nodes()
+        assert a.flags.aligned
+        assert a.flags.c_contiguous
+        assert a.flags.owndata
+        b = ts_fixture.get_individuals_nodes()
+        assert a is not b
+        assert np.all(a == b)
+
+        # Lifetime
+        a1 = ts_fixture.get_individuals_nodes()
+        a2 = a1.copy()
+        assert a1 is not a2
+        del ts_fixture
+        # Do some memory operations
+        a3 = np.ones(10**6)
+        assert np.all(a1 == a2)
+        del a1
+        # Just do something to touch memory
+        a2[:] = 0
+        assert a3 is not a2
+
 
 class StatsInterfaceMixin:
     """

python/tskit/trees.py

@@ -4123,6 +4123,7 @@ def __init__(self, ll_tree_sequence):
         self._individuals_time = None
         self._individuals_population = None
         self._individuals_location = None
+        self._individuals_nodes = None
         # NOTE: when we've implemented read-only access via the underlying
         # tables we can replace these arrays with reference to the read-only
         # tables here (and remove the low-level boilerplate).
@@ -5814,6 +5815,20 @@ def individuals_metadata(self):
             self._individuals_metadata
         )
 
+    @property
+    def individuals_nodes(self):
+        """
+        Return an array of node IDs for each individual in the tree sequence.
+
+        :return: Array of shape (num_individuals, max_ploidy) containing node IDs.
+            Values of -1 indicate unused slots for individuals with ploidy
+            less than the maximum.
+        :rtype: numpy.ndarray (dtype=np.int32)
+        """
+        if self._individuals_nodes is None:
+            self._individuals_nodes = self._ll_tree_sequence.get_individuals_nodes()
+        return self._individuals_nodes
+
     @property
     def nodes_metadata(self):
         """