WIP - switch to iterators from sources

Author: benjeffery

bio2zarr/plink.py

@@ -17,7 +17,7 @@ def __init__(self, path):
         self.num_samples = len(self.samples)
         self.root_attrs = {}
 
-    def write_alleles_to_buffered_array(self, ba, start, stop):
+    def iter_alleles(self, start, stop, num_alleles):
         ref_field = self.bed.allele_1
         alt_field = self.bed.allele_2
 
@@ -26,35 +26,23 @@ def write_alleles_to_buffered_array(self, ba, start, stop):
             ref_field[start:stop],
             alt_field[start:stop],
         ):
-            j = ba.next_buffer_row()
-            ba.buff[j, :] = constants.STR_FILL
-            ba.buff[j, 0] = ref
-            ba.buff[j, 1] = alt
+            alleles = np.full(num_alleles, constants.STR_FILL, dtype="O")
+            alleles[0] = ref
+            alleles[1 : 1 + len(alt)] = alt
+            yield alleles
 
-    def write_other_field_to_buffered_array(self, ba, field_name, start, stop):
+    def iter_field(self, field_name, shape, start, stop):
         data = {
-            "POS": self.bed.bp_position,
+            "position": self.bed.bp_position,
         }[field_name]
         for value in data[start:stop]:
-            j = ba.next_buffer_row()
-            ba.buff[j] = value
-
-    def write_genotypes_to_buffered_array(self, gt, gt_phased, start, stop):
-        bed_chunk = self.bed.read(slice(start, stop), dtype=np.int8).T
-        # Iterate through each sample's genotypes
-        for values in bed_chunk:
-            # Write to genotype array
-            j = gt.next_buffer_row()
-            g = np.zeros_like(gt.buff[j])
-            g[values == -127] = -1  # Missing values
-            g[values == 0] = [1, 1]  # Homozygous ALT (2 in PLINK)
-            g[values == 1] = [1, 0]  # Heterozygous (1 in PLINK)
-            g[values == 2] = [0, 0]  # Homozygous REF (0 in PLINK)
-            gt.buff[j] = g
-
-            # Write to phased array (PLINK data is unphased)
-            j_phased = gt_phased.next_buffer_row()
-            gt_phased.buff[j_phased] = False
+            yield value
+
+    def iter_genotypes(self, start, stop):
+        gt_calls = self.bed.gts.values[start:stop]
+        phased = np.zeros_like(gt_calls, dtype=bool)
+        for idx in range(len(gt_calls)):
+            yield gt_calls[idx], phased[idx]
 
 
 # Import here to avoid circular import
@@ -82,7 +70,7 @@ def generate_schema(
 
     array_specs = [
         schema.ZarrArraySpec.new(
-            vcf_field="POS",
+            vcf_field="position",
             name="variant_position",
             dtype="i4",
             shape=[m],

bio2zarr/vcf2zarr/icf.py

@@ -16,9 +16,127 @@
 from bio2zarr import schema, zarr_utils
 
 from .. import constants, core, provenance, vcf_utils
+from functools import partial
 
 logger = logging.getLogger(__name__)
 
+def sanitise_value_bool(shape, value):
+    x = True
+    if value is None:
+        x = False
+    return x
+
+
+def sanitise_value_float_scalar(shape, value):
+    x = value
+    if value is None:
+        x = [constants.FLOAT32_MISSING]
+    return x[0]
+
+
+def sanitise_value_int_scalar(shape, value):
+    x = value
+    if value is None:
+        x = [constants.INT_MISSING]
+    else:
+        x = sanitise_int_array(value, ndmin=1, dtype=np.int32)
+    return x[0]
+
+
+def sanitise_value_string_scalar(shape, value):
+    if value is None:
+        return "."
+    else:
+        return value[0]
+
+
+def sanitise_value_string_1d(shape, value):
+    if value is None:
+        return np.full(shape, ".", dtype='O')
+    else:
+        value = drop_empty_second_dim(value)
+        result = np.full(shape, "", dtype=value.dtype)
+        result[:value.shape[0]] = value
+        return result
+
+
+def sanitise_value_string_2d(shape, value):
+    if value is None:
+        return np.full(shape, ".", dtype='O')
+    else:
+        result = np.full(shape, "", dtype='O')
+        if value.ndim == 2:
+            result[:value.shape[0], :value.shape[1]] = value
+        else:
+            # Convert 1D array into 2D with appropriate shape
+            for k, val in enumerate(value):
+                result[k, :len(val)] = val
+        return result
+
+
+def drop_empty_second_dim(value):
+    assert len(value.shape) == 1 or value.shape[1] == 1
+    if len(value.shape) == 2 and value.shape[1] == 1:
+        value = value[..., 0]
+    return value
+
+def sanitise_value_float_1d(shape, value):
+    if value is None:
+        return np.full(shape, constants.FLOAT32_MISSING)
+    else:
+        value = np.array(value, ndmin=1, dtype=np.float32, copy=True)
+        # numpy will map None values to Nan, but we need a
+        # specific NaN
+        value[np.isnan(value)] = constants.FLOAT32_MISSING
+        value = drop_empty_second_dim(value)
+        result = np.full(shape, constants.FLOAT32_FILL, dtype=np.float32)
+        result[:value.shape[0]] = value
+        print(result)
+        return result
+
+def sanitise_value_float_2d(shape, value):
+    if value is None:
+        return np.full(shape, constants.FLOAT32_MISSING)
+    else:
+        value = np.array(value, ndmin=2, dtype=np.float32, copy=True)
+        result = np.full(shape, constants.FLOAT32_FILL, dtype=np.float32)
+        result[:, :value.shape[1]] = value
+        print(result)
+        return result
+
+
+def sanitise_int_array(value, ndmin, dtype):
+    if isinstance(value, tuple):
+        value = [
+            constants.VCF_INT_MISSING if x is None else x for x in value
+        ]  # NEEDS TEST
+    value = np.array(value, ndmin=ndmin, copy=True)
+    value[value == constants.VCF_INT_MISSING] = -1
+    value[value == constants.VCF_INT_FILL] = -2
+    # TODO watch out for clipping here!
+    return value.astype(dtype)
+
+
+def sanitise_value_int_1d(shape, value):
+    if value is None:
+        return np.full(shape, -1)
+    else:
+        value = sanitise_int_array(value, 1, np.int32)
+        value = drop_empty_second_dim(value)
+        result = np.full(shape, -2, dtype=np.int32)
+        result[:value.shape[0]] = value
+        return result
+
+
+def sanitise_value_int_2d(shape, value):
+    if value is None:
+        return np.full(shape, -1)
+    else:
+        value = sanitise_int_array(value, 2, np.int32)
+        result = np.full(shape, -2, dtype=np.int32)
+        result[:, :value.shape[1]] = value
+        return result
+
 
 @dataclasses.dataclass
 class VcfFieldSummary(core.JsonDataclass):
@@ -572,35 +690,41 @@ def values(self):
 
     def sanitiser_factory(self, shape):
         """
-        Return a function that sanitised values from this column
-        and writes into a buffer of the specified shape.
+        Return a function that sanitises values from this column
+        and returns a properly formatted array with the specified shape.
+        
+        Args:
+            shape: The shape of the target buffer, used to determine how to format the output
+            
+        Returns:
+            A function that takes a value and returns a sanitised version
         """
-        assert len(shape) <= 3
+        assert len(shape) <= 2
         if self.vcf_field.vcf_type == "Flag":
-            assert len(shape) == 1
-            return zarr_utils.sanitise_value_bool
+            assert len(shape) == 0
+            return partial(sanitise_value_bool, shape)
         elif self.vcf_field.vcf_type == "Float":
-            if len(shape) == 1:
-                return zarr_utils.sanitise_value_float_scalar
-            elif len(shape) == 2:
-                return zarr_utils.sanitise_value_float_1d
+            if len(shape) == 0:
+                return partial(sanitise_value_float_scalar, shape)
+            elif len(shape) == 1:
+                return partial(sanitise_value_float_1d, shape)
             else:
-                return zarr_utils.sanitise_value_float_2d
+                return partial(sanitise_value_float_2d, shape)
         elif self.vcf_field.vcf_type == "Integer":
-            if len(shape) == 1:
-                return zarr_utils.sanitise_value_int_scalar
-            elif len(shape) == 2:
-                return zarr_utils.sanitise_value_int_1d
+            if len(shape) == 0:
+                return partial(sanitise_value_int_scalar, shape)
+            elif len(shape) == 1:
+                return partial(sanitise_value_int_1d, shape)
             else:
-                return zarr_utils.sanitise_value_int_2d
+                return partial(sanitise_value_int_2d, shape)
         else:
             assert self.vcf_field.vcf_type in ("String", "Character")
-            if len(shape) == 1:
-                return zarr_utils.sanitise_value_string_scalar
-            elif len(shape) == 2:
-                return zarr_utils.sanitise_value_string_1d
+            if len(shape) == 0:
+                return partial(sanitise_value_string_scalar, shape)
+            elif len(shape) == 1:
+                return partial(sanitise_value_string_1d, shape)
             else:
-                return zarr_utils.sanitise_value_string_2d
+                return partial(sanitise_value_string_2d, shape)
 
 
 @dataclasses.dataclass
@@ -790,40 +914,33 @@ def root_attrs(self):
             "vcf_header": self.vcf_header,
         }
 
-    def write_alleles_to_buffered_array(self, ba, start, stop):
+    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),
         ):
-            j = ba.next_buffer_row()
-            ba.buff[j, :] = constants.STR_FILL
-            ba.buff[j, 0] = ref[0]
-            ba.buff[j, 1 : 1 + len(alt)] = alt
+            alleles = np.full(num_alleles, constants.STR_FILL, dtype="O")
+            alleles[0] = ref[0]
+            alleles[1 : 1 + len(alt)] = alt
+            yield alleles
 
-    def write_other_field_to_buffered_array(self, ba, field_name, start, stop):
+    def iter_field(self, field_name, shape, start, stop):
         source_field = self.fields[field_name]
-        sanitiser = source_field.sanitiser_factory(ba.buff.shape)
-
+        sanitiser = source_field.sanitiser_factory(shape)
         for value in source_field.iter_values(start, stop):
-            # We write directly into the buffer in the sanitiser function
-            # to make it easier to reason about dimension padding
-            j = ba.next_buffer_row()
-            sanitiser(ba.buff, j, value)
+            yield sanitiser(value)
 
-    def write_genotypes_to_buffered_array(self, gt, gt_phased, start, stop):
+    def iter_genotypes(self, shape, start, stop):
         source_field = self.fields["FORMAT/GT"]
         for value in source_field.iter_values(start, stop):
-            j = gt.next_buffer_row()
-            zarr_utils.sanitise_value_int_2d(
-                gt.buff, j, value[:, :-1] if value is not None else None
-            )
-            j = gt_phased.next_buffer_row()
-            zarr_utils.sanitise_value_int_1d(
-                gt_phased.buff, j, value[:, -1] if value is not None else None
-            )
+            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
 
 
 @dataclasses.dataclass

bio2zarr/writer.py

@@ -87,10 +87,10 @@ class LocalisableFieldDescriptor:
 
 localisable_fields = [
     LocalisableFieldDescriptor(
-        "call_LAD", "FORMAT/AD", zarr_utils.sanitise_int_array, compute_lad_field
+        "call_LAD", "FORMAT/AD", icf.sanitise_int_array, compute_lad_field
     ),
     LocalisableFieldDescriptor(
-        "call_LPL", "FORMAT/PL", zarr_utils.sanitise_int_array, compute_lpl_field
+        "call_LPL", "FORMAT/PL", icf.sanitise_int_array, compute_lpl_field
     ),
 ]
 
@@ -419,24 +419,33 @@ def finalise_partition_array(self, partition_index, buffered_array):
     def encode_array_partition(self, array_spec, partition_index):
         partition = self.metadata.partitions[partition_index]
         ba = self.init_partition_array(partition_index, array_spec.name)
-        self.source.write_other_field_to_buffered_array(
-            ba,
+        for value in self.source.iter_field(
             array_spec.vcf_field,
+            ba.buff.shape[1:],
             partition.start,
             partition.stop,
-        )
+        ):
+            j = ba.next_buffer_row()
+            ba.buff[j] = value
+            
         self.finalise_partition_array(partition_index, ba)
 
     def encode_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")
-        self.source.write_genotypes_to_buffered_array(
-            gt,
-            gt_phased,
+        
+        for genotype, phased in self.source.iter_genotypes(
+            gt.buff.shape[1:],
             partition.start,
-            partition.stop,
-        )
+            partition.stop
+        ):
+            j = gt.next_buffer_row()
+            gt.buff[j] = genotype
+            
+            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)
 
@@ -504,9 +513,9 @@ def encode_alleles_partition(self, partition_index):
         alleles = self.init_partition_array(partition_index, "variant_allele")
         partition = self.metadata.partitions[partition_index]
 
-        self.source.write_alleles_to_buffered_array(
-            alleles, partition.start, partition.stop
-        )
+        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)