Merge pull request #113 from jeromekelleher/array-niceties

Add named dimensions for vcf numbers

Author: jeromekelleher

bio2zarr/vcf.py

@@ -1388,7 +1388,16 @@ def fixed_field_spec(
             # TODO make an option to add in the empty extra dimension
             if field.summary.max_number > 1:
                 shape.append(field.summary.max_number)
-                dimensions.append(field.name)
+                # TODO we should really be checking this to see if the named dimensions
+                # are actually correct.
+                if field.vcf_number == "R":
+                    dimensions.append("alleles")
+                elif field.vcf_number == "A":
+                    dimensions.append("alt_alleles")
+                elif field.vcf_number == "G":
+                    dimensions.append("genotypes")
+                else:
+                    dimensions.append(f"{field.category}_{field.name}_dim")
             variable_name = prefix + field.name
             colspec = ZarrColumnSpec(
                 vcf_field=field.full_name,

tests/test_icf.py

@@ -195,30 +195,34 @@ def schema(self, icf):
         return vcf.VcfZarrSchema.generate(icf)
 
     @pytest.mark.parametrize(
-        ("name", "dtype", "shape"),
+        ("name", "dtype", "shape", "dimensions"),
         [
-            ("variant_II1", "i1", (208,)),
-            ("variant_II2", "i2", (208, 2)),
-            ("variant_IIA", "i2", (208, 2)),
-            ("variant_IIR", "i2", (208, 3)),
-            ("variant_IID", "i2", (208, 7)),
-            ("variant_IF1", "f4", (208,)),
-            ("variant_IF2", "f4", (208, 2)),
-            ("variant_IFA", "f4", (208, 2)),
-            ("variant_IFR", "f4", (208, 3)),
-            ("variant_IFD", "f4", (208, 9)),
-            ("variant_IC1", "U1", (208,)),
-            ("variant_IC2", "U1", (208, 2)),
-            ("variant_IS1", "O", (208,)),
-            ("variant_IS2", "O", (208, 2)),
-            ("call_FS2", "O", (208, 2, 2)),
-            ("call_FC2", "U1", (208, 2, 2)),
+            ("variant_II1", "i1", (208,), ("variants",)),
+            ("variant_II2", "i2", (208, 2), ("variants", "INFO_II2_dim")),
+            ("variant_IIA", "i2", (208, 2), ("variants", "alt_alleles")),
+            ("variant_IIR", "i2", (208, 3), ("variants", "alleles")),
+            ("variant_IID", "i2", (208, 7), ("variants", "INFO_IID_dim")),
+            ("variant_IF1", "f4", (208,), ("variants",)),
+            ("variant_IF2", "f4", (208, 2), ("variants", "INFO_IF2_dim")),
+            ("variant_IFA", "f4", (208, 2), ("variants", "alt_alleles")),
+            ("variant_IFR", "f4", (208, 3), ("variants", "alleles")),
+            ("variant_IFD", "f4", (208, 9), ("variants", "INFO_IFD_dim")),
+            ("variant_IC1", "U1", (208,), ("variants",)),
+            ("variant_IC2", "U1", (208, 2), ("variants", "INFO_IC2_dim")),
+            ("variant_IS1", "O", (208,), ("variants",)),
+            ("variant_IS2", "O", (208, 2), ("variants", "INFO_IS2_dim")),
+            ("call_FS2", "O", (208, 2, 2), ("variants", "samples", "FORMAT_FS2_dim")),
+            ("call_FC2", "U1", (208, 2, 2), ("variants", "samples", "FORMAT_FC2_dim")),
+            ("call_FIG", "i2", (208, 2, 6), ("variants", "samples", "genotypes")),
+            ("call_FIA", "i2", (208, 2, 2), ("variants", "samples", "alt_alleles")),
+            ("call_FIR", "i2", (208, 2, 3), ("variants", "samples", "alleles")),
         ],
     )
-    def test_info_schemas(self, schema, name, dtype, shape):
+    def test_info_schemas(self, schema, name, dtype, shape, dimensions):
         v = schema.columns[name]
         assert v.dtype == dtype
         assert tuple(v.shape) == shape
+        assert v.dimensions == dimensions
 
     def test_info_string1(self, icf):
         non_missing = [v for v in icf["INFO/IS1"].values if v is not None]

tests/test_vcf_examples.py

@@ -384,6 +384,23 @@ def test_missing_contig_vcf(self, ds, tmp_path, path):
             drop_vars = ["contig_id", "variant_contig"]
             xt.assert_equal(ds_c1.drop_vars(drop_vars), ds_c2.drop_vars(drop_vars))
 
+    def test_vcf_dimensions(self, ds):
+        assert ds.call_genotype.dims == ("variants", "samples", "ploidy")
+        assert ds.call_genotype_mask.dims == ("variants", "samples", "ploidy")
+        assert ds.call_genotype_phased.dims == ("variants", "samples")
+        assert ds.call_HQ.dims == ("variants", "samples", "FORMAT_HQ_dim")
+        assert ds.call_DP.dims == ("variants", "samples")
+        assert ds.call_GQ.dims == ("variants", "samples")
+        assert ds.variant_AA.dims == ("variants",)
+        assert ds.variant_NS.dims == ("variants",)
+        assert ds.variant_AN.dims == ("variants",)
+        assert ds.variant_AC.dims == ("variants", "INFO_AC_dim")
+        assert ds.variant_AF.dims == ("variants", "INFO_AF_dim")
+        assert ds.variant_DP.dims == ("variants", )
+        assert ds.variant_DB.dims == ("variants", )
+        assert ds.variant_H2.dims == ("variants", )
+        assert ds.variant_position.dims == ("variants", )
+
 
 class Test1000G2020Example:
     data_path = "tests/data/vcf/1kg_2020_chrM.vcf.gz"