Move individual nodes to the caller, use numpy assign

Author: benjeffery

bio2zarr/tskit.py

@@ -9,79 +9,46 @@
 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, isolated_as_missing=False):
+    def __init__(
+        self,
+        ts_path,
+        individual_nodes,
+        sample_ids=None,
+        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()
         self.positions = self.ts.sites_position
 
+        self._num_samples = individual_nodes.shape[0]
+        if self._num_samples < 1:
+            raise ValueError("individual_nodes must have at least one sample")
+        self.max_ploidy = individual_nodes.shape[1]
+        if sample_ids is None:
+            sample_ids = [f"tsk_{j}" for j in range(self._num_samples)]
+        elif len(sample_ids) != self._num_samples:
+            raise ValueError(
+                f"Length of sample_ids ({len(sample_ids)}) does not match "
+                f"number of samples ({self._num_samples})"
+            )
+
+        self._samples = [vcz.Sample(id=sample_id) for sample_id in sample_ids]
+
+        self.tskit_samples = np.unique(individual_nodes[individual_nodes >= 0])
+        if len(self.tskit_samples) < 1:
+            raise ValueError("individual_nodes must have at least one valid sample")
+        node_id_to_index = {node_id: i for i, node_id in enumerate(self.tskit_samples)}
+        valid_mask = individual_nodes >= 0
+        self.sample_indices, self.ploidy_indices = np.where(valid_mask)
+        self.genotype_indices = np.array(
+            [node_id_to_index[node_id] for node_id in individual_nodes[valid_mask]]
+        )
+
     @property
     def path(self):
         return self._path
@@ -106,21 +73,6 @@ def root_attrs(self):
     def contigs(self):
         return [vcz.Contig(id=self.contig_id)]
 
-    def _make_sample_mapping(self):
-        ts = self.ts
-        
-        # 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):
         yield from (0 for _ in range(start, stop))
 
@@ -139,6 +91,7 @@ def iter_alleles_and_genotypes(self, start, stop, shape, num_alleles):
             isolated_as_missing=self.isolated_as_missing,
             left=self.positions[start],
             right=self.positions[stop] if stop < self.num_records else None,
+            samples=self.tskit_samples,
         ):
             gt = np.full(shape, constants.INT_FILL, dtype=np.int8)
             alleles = np.full(num_alleles, constants.STR_FILL, dtype="O")
@@ -149,12 +102,9 @@ def iter_alleles_and_genotypes(self, start, stop, shape, num_alleles):
                 assert i < num_alleles
                 alleles[i] = allele
 
-            # 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]
+            gt[self.sample_indices, self.ploidy_indices] = variant.genotypes[
+                self.genotype_indices
+            ]
 
             yield alleles, (gt, phased)
 
@@ -253,7 +203,9 @@ def generate_schema(
 def convert(
     ts_path,
     zarr_path,
+    individual_nodes,
     *,
+    sample_ids=None,
     contig_id=None,
     isolated_as_missing=False,
     variants_chunk_size=None,
@@ -263,6 +215,8 @@ def convert(
 ):
     tskit_format = TskitFormat(
         ts_path,
+        individual_nodes,
+        sample_ids=sample_ids,
         contig_id=contig_id,
         isolated_as_missing=isolated_as_missing,
     )