Add genotypes dimension to schema

Author: tomwhite

bio2zarr/icf.py

@@ -1050,12 +1050,32 @@ def generate_schema(
             "samples": vcz.VcfZarrDimension(
                 size=n, chunk_size=samples_chunk_size or 10000
             ),
-            # ploidy added conditionally below
+            # ploidy and genotypes added conditionally below
             "alleles": vcz.VcfZarrDimension(size=max_alleles),
             "alt_alleles": vcz.VcfZarrDimension(size=max_alleles - 1),
             "filters": vcz.VcfZarrDimension(size=self.metadata.num_filters),
         }
 
+        # Add ploidy and genotypes dimensions only when needed
+        gt_field = None
+        max_genotypes = 0
+        for field in self.metadata.format_fields:
+            if field.name == "GT":
+                gt_field = field
+                continue
+            if field.vcf_number == "G":
+                max_genotypes = max(max_genotypes, field.summary.max_number)
+        if gt_field is not None:
+            ploidy = max(gt_field.summary.max_number - 1, 1)
+            dimensions["ploidy"] = vcz.VcfZarrDimension(size=ploidy)
+            max_genotypes = math.comb(max_alleles + ploidy - 1, ploidy)
+            dimensions["genotypes"] = vcz.VcfZarrDimension(size=max_genotypes)
+        else:
+            if max_genotypes > 0:
+                # there is no GT field, but there is at least one Number=G field,
+                # so need to define genotypes dimension
+                dimensions["genotypes"] = vcz.VcfZarrDimension(size=max_genotypes)
+
         schema_instance = vcz.VcfZarrSchema(
             format_version=vcz.ZARR_SCHEMA_FORMAT_VERSION,
             dimensions=dimensions,
@@ -1128,18 +1148,12 @@ def fixed_field_spec(name, dtype, source=None, dimensions=("variants",)):
             [spec_from_field(field) for field in self.metadata.info_fields]
         )
 
-        gt_field = None
         for field in self.metadata.format_fields:
             if field.name == "GT":
-                gt_field = field
                 continue
             array_specs.append(spec_from_field(field))
 
         if gt_field is not None and n > 0:
-            ploidy = max(gt_field.summary.max_number - 1, 1)
-            # Add ploidy dimension only when needed
-            schema_instance.dimensions["ploidy"] = vcz.VcfZarrDimension(size=ploidy)
-
             array_specs.append(
                 vcz.ZarrArraySpec(
                     name="call_genotype_phased",

bio2zarr/vcz.py

@@ -190,9 +190,15 @@ def from_field(
                 )
             if max_alt_alleles > 0:
                 dimensions.append("alt_alleles")
-        elif max_number > 0 and vcf_field.vcf_number == "G":
-            # TODO: need max_genotypes
-            dimensions.append("genotypes")
+        elif vcf_field.vcf_number == "G":
+            max_genotypes = schema.dimensions["genotypes"].size
+            if max_number > max_genotypes:
+                raise ValueError(
+                    f"Max number of values {max_number} exceeds max genotypes "
+                    f"{max_genotypes} for {vcf_field.full_name}"
+                )
+            if max_genotypes > 0:
+                dimensions.append("genotypes")
         elif max_number > 1 or vcf_field.full_name == "FORMAT/LAA":
             dimensions.append(f"{vcf_field.category}_{vcf_field.name}_dim")
         if dimensions[-1] not in schema.dimensions: