Iterate alleles and genotypes at the same time

Author: benjeffery

bio2zarr/icf.py

@@ -970,19 +970,6 @@ def root_attrs(self):
             "vcf_header": self.vcf_header,
         }
 
-    def iter_alleles(self, start, stop, num_alleles):
-        ref_field = self.fields["REF"]
-        alt_field = self.fields["ALT"]
-
-        for ref, alt in zip(
-            ref_field.iter_values(start, stop),
-            alt_field.iter_values(start, stop),
-        ):
-            alleles = np.full(num_alleles, constants.STR_FILL, dtype="O")
-            alleles[0] = ref[0]
-            alleles[1 : 1 + len(alt)] = alt
-            yield alleles
-
     def iter_id(self, start, stop):
         for value in self.fields["ID"].iter_values(start, stop):
             if value is not None:
@@ -1025,14 +1012,32 @@ def iter_field(self, field_name, shape, start, stop):
         for value in source_field.iter_values(start, stop):
             yield sanitiser(value)
 
-    def iter_genotypes(self, shape, start, stop):
-        source_field = self.fields["FORMAT/GT"]
-        for value in source_field.iter_values(start, stop):
-            genotypes = value[:, :-1] if value is not None else None
-            phased = value[:, -1] if value is not None else None
-            sanitised_genotypes = sanitise_value_int_2d(shape, genotypes)
-            sanitised_phased = sanitise_value_int_1d(shape[:-1], phased)
-            yield sanitised_genotypes, sanitised_phased
+    def iter_alleles_and_genotypes(
+        self, start, stop, include_genotypes, shape, num_alleles
+    ):
+        ref_field = self.fields["REF"]
+        alt_field = self.fields["ALT"]
+        # Create a dummy gt iterator if genotypes are not included
+
+        for ref, alt, gt in zip(
+            ref_field.iter_values(start, stop),
+            alt_field.iter_values(start, stop),
+            (None for _ in range(stop - start))
+            if not include_genotypes
+            else self.fields["FORMAT/GT"].iter_values(start, stop),
+        ):
+            alleles = np.full(num_alleles, constants.STR_FILL, dtype="O")
+            alleles[0] = ref[0]
+            alleles[1 : 1 + len(alt)] = alt
+
+            if include_genotypes:
+                genotypes = gt[:, :-1] if gt is not None else None
+                phased = gt[:, -1] if gt is not None else None
+                sanitised_genotypes = sanitise_value_int_2d(shape, genotypes)
+                sanitised_phased = sanitise_value_int_1d(shape[:-1], phased)
+                yield alleles, sanitised_genotypes, sanitised_phased
+            else:
+                yield alleles, None, None
 
     def generate_schema(
         self, variants_chunk_size=None, samples_chunk_size=None, local_alleles=None

bio2zarr/plink.py

@@ -31,34 +31,32 @@ def samples(self):
     def num_samples(self):
         return len(self.samples)
 
-    def iter_alleles(self, start, stop, num_alleles):
-        ref_field = self.bed.allele_1
-        alt_field = self.bed.allele_2
-
-        for ref, alt in zip(
-            ref_field[start:stop],
-            alt_field[start:stop],
-        ):
-            alleles = np.full(num_alleles, constants.STR_FILL, dtype="O")
-            alleles[0] = ref
-            alleles[1 : 1 + len(alt)] = alt
-            yield alleles
-
     def iter_field(self, field_name, shape, start, stop):
         assert field_name == "position"  # Only position field is supported from plink
         yield from self.bed.bp_position[start:stop]
 
-    def iter_genotypes(self, shape, start, stop):
+    def iter_alleles_and_genotypes(
+        self, start, stop, include_genotypes, shape, num_alleles
+    ):
+        if not include_genotypes:
+            raise NotImplementedError("include_genotypes must be True for plink format")
+        ref_field = self.bed.allele_1
+        alt_field = self.bed.allele_2
         bed_chunk = self.bed.read(slice(start, stop), dtype=np.int8).T
         gt = np.zeros(shape, dtype=np.int8)
         phased = np.zeros(shape[:-1], dtype=bool)
-        for values in bed_chunk:
+        for i, (ref, alt) in enumerate(
+            zip(ref_field[start:stop], alt_field[start:stop])
+        ):
+            alleles = np.full(num_alleles, constants.STR_FILL, dtype="O")
+            alleles[0] = ref
+            alleles[1 : 1 + len(alt)] = alt
             gt[:] = 0
-            gt[values == -127] = -1
-            gt[values == 2] = 1
-            gt[values == 1, 0] = 1
+            gt[bed_chunk[i] == -127] = -1
+            gt[bed_chunk[i] == 2] = 1
+            gt[bed_chunk[i] == 1, 0] = 1
 
-            yield gt, phased
+            yield alleles, gt, phased
 
     def generate_schema(
         self,

bio2zarr/vcz.py

@@ -69,11 +69,9 @@ def root_attrs(self):
         return {}
 
     @abc.abstractmethod
-    def iter_alleles(self, start, stop, num_alleles):
-        pass
-
-    @abc.abstractmethod
-    def iter_genotypes(self, start, stop, num_alleles):
+    def iter_alleles_and_genotypes(
+        self, start, stop, include_genotypes, shape, num_alleles
+    ):
         pass
 
     def iter_id(self, start, stop):
@@ -718,12 +716,11 @@ def encode_partition(self, partition_index):
             self.encode_filters_partition(partition_index)
         if "variant_contig" in all_field_names:
             self.encode_contig_partition(partition_index)
-        self.encode_alleles_partition(partition_index)
+        self.encode_alleles_and_genotypes_partition(partition_index)
         for array_spec in self.schema.fields:
             if array_spec.source is not None:
                 self.encode_array_partition(array_spec, partition_index)
         if self.has_genotypes():
-            self.encode_genotypes_partition(partition_index)
             self.encode_genotype_mask_partition(partition_index)
         if self.has_local_alleles():
             self.encode_local_alleles_partition(partition_index)
@@ -774,22 +771,33 @@ def encode_array_partition(self, array_spec, partition_index):
 
         self.finalise_partition_array(partition_index, ba)
 
-    def encode_genotypes_partition(self, partition_index):
+    def encode_alleles_and_genotypes_partition(self, partition_index):
         partition = self.metadata.partitions[partition_index]
-        gt = self.init_partition_array(partition_index, "call_genotype")
-        gt_phased = self.init_partition_array(partition_index, "call_genotype_phased")
+        alleles = self.init_partition_array(partition_index, "variant_allele")
+        has_gt = self.has_genotypes()
+        shape = None
+        if has_gt:
+            gt = self.init_partition_array(partition_index, "call_genotype")
+            gt_phased = self.init_partition_array(
+                partition_index, "call_genotype_phased"
+            )
+            shape = gt.buff.shape[1:]
 
-        for genotype, phased in self.source.iter_genotypes(
-            gt.buff.shape[1:], partition.start, partition.stop
+        for alleles_value, genotype, phased in self.source.iter_alleles_and_genotypes(
+            partition.start, partition.stop, has_gt, shape, alleles.array.shape[1]
         ):
-            j = gt.next_buffer_row()
-            gt.buff[j] = genotype
+            j_alleles = alleles.next_buffer_row()
+            alleles.buff[j_alleles] = alleles_value
+            if has_gt:
+                j = gt.next_buffer_row()
+                gt.buff[j] = genotype
+                j_phased = gt_phased.next_buffer_row()
+                gt_phased.buff[j_phased] = phased
 
-            j_phased = gt_phased.next_buffer_row()
-            gt_phased.buff[j_phased] = phased
-
-        self.finalise_partition_array(partition_index, gt)
-        self.finalise_partition_array(partition_index, gt_phased)
+        self.finalise_partition_array(partition_index, alleles)
+        if has_gt:
+            self.finalise_partition_array(partition_index, gt)
+            self.finalise_partition_array(partition_index, gt_phased)
 
     def encode_genotype_mask_partition(self, partition_index):
         partition = self.metadata.partitions[partition_index]
@@ -851,18 +859,6 @@ def encode_local_allele_fields_partition(self, partition_index):
                 buff.buff[j] = descriptor.convert(value, la)
             self.finalise_partition_array(partition_index, buff)
 
-    def encode_alleles_partition(self, partition_index):
-        alleles = self.init_partition_array(partition_index, "variant_allele")
-        partition = self.metadata.partitions[partition_index]
-
-        for value in self.source.iter_alleles(
-            partition.start, partition.stop, alleles.array.shape[1]
-        ):
-            j = alleles.next_buffer_row()
-            alleles.buff[j] = value
-
-        self.finalise_partition_array(partition_index, alleles)
-
     def encode_id_partition(self, partition_index):
         vid = self.init_partition_array(partition_index, "variant_id")
         vid_mask = self.init_partition_array(partition_index, "variant_id_mask")