Simplify and use individuals_node method

Author: benjeffery

bio2zarr/tskit.py

@@ -9,14 +9,77 @@
 logger = logging.getLogger(__name__)
 
 
+import numpy as np
+import tskit
+
+def individual_nodes(ts):
+    """
+    Convert a tree sequence with individuals to a 2D array of node IDs.
+    
+    Parameters
+    ----------
+    ts : tskit.TreeSequence
+        The tree sequence to convert
+        
+    Returns
+    -------
+    numpy.ndarray
+        Array of shape (num_individuals, max_ploidy) containing node IDs.
+        Values of -1 indicate unused slots for individuals with ploidy
+        less than the maximum.
+        
+    Raises
+    ------
+    ValueError
+        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 
+        individuals has no samples.
+    """
+    if ts.num_individuals == 0:
+        raise ValueError("Tree sequence has no individuals")
+    
+    individuals = np.unique(ts.nodes_individual[ts.samples()])
+    if len(individuals) == 1 and individuals[0] == tskit.NULL:
+        raise ValueError("No samples refer to individuals")
+    
+    # np.unique sorts the argument, so if NULL (-1) is present it will be first
+    if individuals[0] == tskit.NULL:
+        raise ValueError(
+            "Sample nodes must all be associated with individuals"
+        )
+    
+    max_ploidy = 0
+    for i in range(ts.num_individuals):
+        ind = ts.individual(i)
+        max_ploidy = max(max_ploidy, len(ind.nodes))
+    
+    # Initialize output array with -1 (indicating no node)
+    result = np.full((ts.num_individuals, max_ploidy), -1, dtype=np.int32)
+    
+    for i in range(ts.num_individuals):
+        ind = ts.individual(i)
+        if len(ind.nodes) == 0:
+            raise ValueError(f"Individual {i} not associated with any nodes")
+        
+        is_sample = {ts.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
+
 class TskitFormat(vcz.Source):
-    def __init__(self, ts_path, contig_id=None, ploidy=None, isolated_as_missing=False):
+    def __init__(self, ts_path, contig_id=None, isolated_as_missing=False):
         self._path = ts_path
         self.ts = tskit.load(ts_path)
         self.contig_id = contig_id if contig_id is not None else "1"
         self.isolated_as_missing = isolated_as_missing
 
-        self._make_sample_mapping(ploidy)
+        self._make_sample_mapping()
         self.positions = self.ts.sites_position
 
     @property
@@ -43,62 +106,19 @@ def root_attrs(self):
     def contigs(self):
         return [vcz.Contig(id=self.contig_id)]
 
-    def _make_sample_mapping(self, ploidy):
+    def _make_sample_mapping(self):
         ts = self.ts
-        self.individual_ploidies = []
-        self.max_ploidy = 0
-
-        if ts.num_individuals > 0 and ploidy is not None:
-            raise ValueError(
-                "Cannot specify ploidy when individuals are present in tables"
-            )
-
-        # Find all sample nodes that reference individuals
-        individuals = np.unique(ts.nodes_individual[ts.samples()])
-        if len(individuals) == 1 and individuals[0] == tskit.NULL:
-            # No samples refer to individuals
-            individuals = None
-        else:
-            # np.unique sorts the argument, so if NULL (-1) is present it
-            # will be the first value.
-            if individuals[0] == tskit.NULL:
-                raise ValueError(
-                    "Sample nodes must either all be associated with individuals "
-                    "or not associated with any individuals"
-                )
-
-        if individuals is not None:
-            self.sample_ids = []
-            for i in individuals:
-                if i < 0 or i >= self.ts.num_individuals:
-                    raise ValueError("Invalid individual IDs provided.")
-                ind = self.ts.individual(i)
-                if len(ind.nodes) == 0:
-                    raise ValueError(f"Individual {i} not associated with a node")
-                is_sample = {ts.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"
-                    )
-                self.sample_ids.extend(ind.nodes)
-                self.individual_ploidies.append(len(ind.nodes))
-                self.max_ploidy = max(self.max_ploidy, len(ind.nodes))
-        else:
-            if ploidy is None:
-                ploidy = 1
-            if ploidy < 1:
-                raise ValueError("Ploidy must be >= 1")
-            if ts.num_samples % ploidy != 0:
-                raise ValueError("Sample size must be divisible by ploidy")
-            self.individual_ploidies = np.full(
-                ts.num_samples // ploidy, ploidy, dtype=np.int32
-            )
-            self.max_ploidy = ploidy
-            self.sample_ids = (ts.nodes_flags & tskit.NODE_IS_SAMPLE).nonzero()[0]
-
-        self._num_samples = len(self.individual_ploidies)
-
+        
+        # Use individual_nodes to get the mapping between individuals and nodes
+        try:
+            # Get a 2D array of node IDs for each individual
+            self.node_ids_array = individual_nodes(ts)
+            self._num_samples = ts.num_individuals
+            self.max_ploidy = self.node_ids_array.shape[1]
+            
+        except ValueError as e:
+            raise ValueError(f"Error mapping individuals to nodes: {e}") from e
+        
         self._samples = [vcz.Sample(id=f"tsk_{j}") for j in range(self.num_samples)]
 
     def iter_contig(self, start, stop):
@@ -111,25 +131,11 @@ def iter_field(self, field_name, shape, start, stop):
         else:
             raise ValueError(f"Unknown field {field_name}")
 
-    def iter_alleles(self, start, stop, num_alleles):
-        for variant in self.ts.variants(
-            samples=self.sample_ids,
-            isolated_as_missing=self.isolated_as_missing,
-            left=self.positions[start],
-            right=self.positions[stop] if stop < self.num_records else None,
-        ):
-            alleles = np.full(num_alleles, constants.STR_FILL, dtype="O")
-            for i, allele in enumerate(variant.alleles):
-                assert i < num_alleles
-                alleles[i] = allele
-            yield alleles
-
     def iter_alleles_and_genotypes(self, start, stop, shape, num_alleles):
-        # In tskit, all genotypes are considered phased
+        # All genotypes in tskit are considered phased
         phased = np.ones(shape[:-1], dtype=bool)
 
         for variant in self.ts.variants(
-            samples=self.sample_ids,
             isolated_as_missing=self.isolated_as_missing,
             left=self.positions[start],
             right=self.positions[stop] if stop < self.num_records else None,
@@ -143,13 +149,12 @@ def iter_alleles_and_genotypes(self, start, stop, shape, num_alleles):
                 assert i < num_alleles
                 alleles[i] = allele
 
-            genotypes = variant.genotypes
-            sample_index = 0
-            for i, ploidy in enumerate(self.individual_ploidies):
-                for j in range(ploidy):
-                    if j < self.max_ploidy:
-                        gt[i, j] = genotypes[sample_index + j]
-                sample_index += ploidy
+            # For each individual, get genotypes for their nodes
+            for i in range(self.num_samples):
+                for j in range(self.max_ploidy):
+                    node_id = self.node_ids_array[i, j]
+                    if node_id >= 0:  # Skip -1 entries (unused slots)
+                        gt[i, j] = variant.genotypes[node_id]
 
             yield alleles, (gt, phased)
 
@@ -250,7 +255,6 @@ def convert(
     zarr_path,
     *,
     contig_id=None,
-    ploidy=None,
     isolated_as_missing=False,
     variants_chunk_size=None,
     samples_chunk_size=None,
@@ -260,7 +264,6 @@ def convert(
     tskit_format = TskitFormat(
         ts_path,
         contig_id=contig_id,
-        ploidy=ploidy,
         isolated_as_missing=isolated_as_missing,
     )
     schema_instance = tskit_format.generate_schema(

tests/test_ts.py

@@ -154,32 +154,12 @@ def test_sample_mapping_with_individuals(self, simple_ts):
         assert format_obj.max_ploidy == 2
         assert format_obj.individual_ploidies == [2, 2]
 
-        # Should raise error if ploidy specified with individuals
-        with pytest.raises(
-            ValueError, match="Cannot specify ploidy when individuals are present"
-        ):
-            ts.TskitFormat(ts_path, ploidy=2)
-
-    def test_sample_mapping_without_individuals(self, no_individuals_ts):
-        ts_path, tree_sequence = no_individuals_ts
-
-        # Default ploidy should be 1
-        format_obj = ts.TskitFormat(ts_path)
-        assert format_obj.num_samples == 4
-        assert format_obj.max_ploidy == 1
-        assert list(format_obj.individual_ploidies) == [1, 1, 1, 1]
-
-        # Explicitly set ploidy to 2
-        format_obj = ts.TskitFormat(ts_path, ploidy=2)
-        assert format_obj.num_samples == 2
-        assert format_obj.max_ploidy == 2
-        assert list(format_obj.individual_ploidies) == [2, 2]
-
-        with pytest.raises(ValueError, match="Ploidy must be >= 1"):
-            ts.TskitFormat(ts_path, ploidy=0)
-
-        with pytest.raises(ValueError, match="Sample size must be divisible by ploidy"):
-            ts.TskitFormat(ts_path, ploidy=3)
+    def test_no_individuals(self, no_individuals_ts):
+        """Test that tree sequences without individuals raise an error."""
+        ts_path, _ = no_individuals_ts
+        
+        with pytest.raises(ValueError, match="Tree sequence has no individuals"):
+            ts.TskitFormat(ts_path)
 
     def test_schema_generation(self, simple_ts):
         ts_path, _ = simple_ts
@@ -307,8 +287,7 @@ def test_variable_ploidy(self, tmp_path):
         format_obj = ts.TskitFormat(ts_path)
 
         assert format_obj.max_ploidy == 3
-        assert format_obj.individual_ploidies == [2, 3]
-
+        
         shape = (2, 3)  # (num_samples, max_ploidy)
         results = list(format_obj.iter_alleles_and_genotypes(0, 2, shape, 2))
 
@@ -387,3 +366,103 @@ def insert_branch_sites(ts, m=1):
         # Individual 2 should have missing values (-1) when isolated_as_missing=True
         expected_gt_missing = np.array([[1], [0], [-1]])
         assert np.array_equal(gt_missing, expected_gt_missing)
+
+
+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)
+        tree_sequence = tables.tree_sequence()
+        
+        result = ts.individual_nodes(tree_sequence)
+        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)
+        
+        tree_sequence = tables.tree_sequence()
+        
+        result = ts.individual_nodes(tree_sequence)
+        
+        assert result.shape == (3, 3)
+        
+        expected = np.array([
+            [0, 1, -1],  # Diploid
+            [2, -1, -1],  # Haploid
+            [3, 4, 5]    # 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)
+        tree_sequence = tables.tree_sequence()
+        
+        with pytest.raises(ValueError, match="Tree sequence has no individuals"):
+            ts.individual_nodes(tree_sequence)
+    
+    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)
+        tree_sequence = tables.tree_sequence()
+        
+        with pytest.raises(ValueError, match="No samples refer to individuals"):
+            ts.individual_nodes(tree_sequence)
+    
+    def test_mixed_individual_references(self):
+        tables = tskit.TableCollection(sequence_length=100)
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")
+        # One node with individual reference
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
+        # One node without individual reference
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        tree_sequence = tables.tree_sequence()
+        
+        with pytest.raises(ValueError, match="Sample nodes must all be associated with individuals"):
+            ts.individual_nodes(tree_sequence)
+    
+    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)
+        tree_sequence = tables.tree_sequence()
+        
+        with pytest.raises(ValueError, match="Individual 1 not associated with any nodes"):
+            ts.individual_nodes(tree_sequence)
+    
+    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)
+        tree_sequence = tables.tree_sequence()
+        
+        with pytest.raises(ValueError, match="has nodes that are sample and non-samples"):
+            ts.individual_nodes(tree_sequence)