Some basic testing of local alleles

Author: jeromekelleher

bio2zarr/vcf2zarr/vcz.py

@@ -204,14 +204,15 @@ def convert_local_allele_field_types(fields):
 
     shape = gt.shape[:-1]
     chunks = gt.chunks[:-1]
+    dimensions = gt.dimensions[:-1]
 
     la = ZarrArraySpec.new(
         vcf_field=None,
         name="call_LA",
         dtype="i1",
         shape=gt.shape,
         chunks=gt.chunks,
-        dimensions=gt.dimensions,  # FIXME
+        dimensions=(*dimensions, "local_alleles"),
         description=(
             "0-based indices into REF+ALT, indicating which alleles"
             " are relevant (local) for the current sample"
@@ -224,8 +225,8 @@ def convert_local_allele_field_types(fields):
         ad.vcf_field = None
         ad.shape = (*shape, 2)
         ad.chunks = (*chunks, 2)
+        ad.dimensions = (*dimensions, "local_alleles")
         ad.description += " (local-alleles)"
-        # TODO fix dimensions
 
     pl = fields_by_name.get("call_PL", None)
     if pl is not None:
@@ -235,7 +236,7 @@ def convert_local_allele_field_types(fields):
         pl.shape = (*shape, 3)
         pl.chunks = (*chunks, 3)
         pl.description += " (local-alleles)"
-        # TODO fix dimensions
+        pl.dimensions = (*dimensions, "local_" + pl.dimensions[-1])
     return [*fields, la]
 
 

tests/test_vcf_examples.py

@@ -716,16 +716,19 @@ def test_call_LAD(self, ds):
             [[446, -2], [393, -2], [486, -2]],
         ]
         nt.assert_array_equal(ds.call_LAD.values, call_LAD)
+        assert ds.call_LAD.dims == ("variants", "samples", "local_alleles")
 
     def test_call_LA(self, ds):
         # All the genotypes are 0/0
         call_LA = np.full((23, 3, 2), -2)
         call_LA[:, :, 0] = 0
         nt.assert_array_equal(ds.call_LA.values, call_LA)
+        assert ds.call_LA.dims == ("variants", "samples", "local_alleles")
 
     def test_call_LPL(self, ds):
         call_LPL = np.tile([0, -2, -2], (23, 3, 1))
         nt.assert_array_equal(ds.call_LPL.values, call_LPL)
+        assert ds.call_LPL.dims == ("variants", "samples", "local_genotypes")
 
 
 class Test1000G2020AnnotationsExample:

tests/test_vcz.py

@@ -40,6 +40,17 @@ def schema(schema_path):
         return vcf2zarr.VcfZarrSchema.fromjson(f.read())
 
 
+@pytest.fixture(scope="module")
+def local_alleles_schema(icf_path, tmp_path_factory):
+    # FIXME: this is stupid way of getting a test fixture, should
+    # be much easier.
+    out = tmp_path_factory.mktemp("data") / "example.schema.json"
+    with open(out, "w") as f:
+        vcf2zarr.mkschema(icf_path, f, local_alleles=True)
+    with open(out) as f:
+        return vcf2zarr.VcfZarrSchema.fromjson(f.read())
+
+
 @pytest.fixture(scope="module")
 def zarr_path(icf_path, tmp_path_factory):
     out = tmp_path_factory.mktemp("data") / "example.zarr"
@@ -445,6 +456,36 @@ def test_call_GQ(self, schema):
         }
 
 
+class TestLocalAllelesDefaultSchema:
+    def test_differences(self, schema, local_alleles_schema):
+        assert len(schema.fields) == len(local_alleles_schema.fields) - 1
+        non_local = [f for f in local_alleles_schema.fields if f.name != "call_LA"]
+        assert schema.fields == non_local
+
+    def test_call_LA(self, local_alleles_schema):
+        d = get_field_dict(local_alleles_schema, "call_LA")
+        assert d == {
+            "name": "call_LA",
+            "dtype": "i1",
+            "shape": (9, 3, 2),
+            "chunks": (1000, 10000, 2),
+            "dimensions": ("variants", "samples", "local_alleles"),
+            "description": (
+                "0-based indices into REF+ALT, indicating which alleles"
+                " are relevant (local) for the current sample"
+            ),
+            "vcf_field": None,
+            "compressor": {
+                "id": "blosc",
+                "cname": "zstd",
+                "clevel": 7,
+                "shuffle": 0,
+                "blocksize": 0,
+            },
+            "filters": tuple(),
+        }
+
+
 class TestVcfDescriptions:
     @pytest.mark.parametrize(
         ("field", "description"),