Add individual nodes method

Author: benjeffery

python/tests/test_highlevel.py

@@ -5491,3 +5491,91 @@ def test_error_if_no_schema(self, table_name):
         ts = msprime.simulate(10)
         with pytest.raises(NotImplementedError):
             getattr(ts, f"{table_name}_metadata")
+
+
+class TestIndividualNodes:
+    def test_basic_individual_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.individual_nodes()
+        assert result.shape == (2, 2)
+        assert np.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.individual_nodes()
+
+        assert result.shape == (3, 3)
+
+        expected = np.array(
+            [[0, 1, -1], [2, -1, -1], [3, 4, 5]]  # Diploid  # Haploid  # Triploid
+        )
+        assert np.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()
+
+        with pytest.raises(ValueError, match="Tree sequence has no individuals"):
+            _ = ts.individual_nodes()
+
+    def test_no_samples_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()
+
+        with pytest.raises(ValueError, match="No nodes refer to individuals"):
+            _ = ts.individual_nodes()
+
+    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()
+
+        with pytest.raises(
+            ValueError, match="Individual 1 not associated with any nodes"
+        ):
+            _ = ts.individual_nodes()
+
+    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()
+        with pytest.raises(
+            ValueError, match="has nodes that are sample and non-samples"
+        ):
+            _ = ts.individual_nodes()

python/tskit/trees.py

@@ -10476,6 +10476,45 @@ def ld_matrix(
             mode=mode,
         )
 
+    def individual_nodes(self):
+        """
+        Return an array of node IDs for each individual in the tree sequence.
+        Errors if the tree sequence has no individuals, if any sample doesn't
+        have an individual, if individuals have nodes that are both samples
+        and non-samples, or if an individual has no samples.
+
+        :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.num_individuals == 0:
+            raise ValueError("Tree sequence has no individuals")
+
+        max_ploidy = 0
+        for i in range(self.num_individuals):
+            ind = self.individual(i)
+            max_ploidy = max(max_ploidy, len(ind.nodes))
+        if max_ploidy == 0:
+            raise ValueError("No nodes refer to individuals")
+
+        # Initialize output array with -1 (indicating no node)
+        result = np.full((self.num_individuals, max_ploidy), -1, dtype=np.int32)
+
+        for i in range(self.num_individuals):
+            ind = self.individual(i)
+            if len(ind.nodes) == 0:
+                raise ValueError(f"Individual {i} not associated with any nodes")
+            is_sample = {self.node(u).is_sample() for u in ind.nodes}
+            if len(is_sample) != 1:
+                raise ValueError(
+                    f"Individual {ind.id} has nodes that are sample and non-samples"
+                )
+            for j, node_id in enumerate(ind.nodes):
+                result[i, j] = node_id
+
+        return result
+
     ############################################
     #
     # Deprecated APIs. These are either already unsupported, or will be unsupported in a