Add round-trip tests using plink's VCF output

Author: jeromekelleher

bio2zarr/cli.py

@@ -561,7 +561,9 @@ def convert_plink(
     samples_chunk_size,
 ):
     """
-    In development; DO NOT USE!
+    Convert plink fileset to VCF Zarr. Results are equivalent to
+    `plink1.9 --bfile prefix --recode vcf-iid --out tmp` then running
+    `vcf2zarr convert tmp.vcf zarr_path`
     """
     setup_logging(verbose)
     plink.convert(

bio2zarr/plink.py

@@ -88,7 +88,8 @@ def iter_alleles_and_genotypes(self, start, stop, shape, num_alleles):
             gt[bed_chunk[i] == 2] = 1
             gt[bed_chunk[i] == 1, 0] = 1
 
-            yield vcz.VariantData(max(len(a) for a in alleles), alleles, gt, phased)
+            # rlen is the length of the REF in PLINK as there's no END annotations
+            yield vcz.VariantData(len(alleles[0]), alleles, gt, phased)
 
     def generate_schema(
         self,
@@ -117,6 +118,9 @@ def generate_schema(
             f"variants={dimensions['variants'].chunk_size}, "
             f"samples={dimensions['samples'].chunk_size}"
         )
+        # If we don't have SVLEN or END annotations, the rlen field is defined
+        # as the length of the REF
+        max_len = self.bed.allele_2.itemsize
 
         array_specs = [
             vcz.ZarrArraySpec(
@@ -147,7 +151,7 @@ def generate_schema(
             vcz.ZarrArraySpec(
                 source=None,
                 name="variant_length",
-                dtype="i4",
+                dtype=core.min_int_dtype(0, max_len),
                 dimensions=["variants"],
                 description="Length of each variant",
             ),

tests/test_cli.py

@@ -1026,23 +1026,6 @@ def test_part_size_multiple_vcfs(self):
         ]
 
 
-@pytest.mark.parametrize(
-    "cmd",
-    [
-        main.bio2zarr,
-        cli.vcf2zarr_main,
-        cli.plink2zarr,
-        cli.vcfpartition,
-        cli.tskit2zarr,
-    ],
-)
-def test_version(cmd):
-    runner = ct.CliRunner()
-    result = runner.invoke(cmd, ["--version"], catch_exceptions=False)
-    s = f"version {provenance.__version__}\n"
-    assert result.stdout.endswith(s)
-
-
 class TestTskitEndToEnd:
     def test_convert(self, tmp_path):
         ts_path = "tests/data/ts/example.trees"
@@ -1060,3 +1043,39 @@ def test_convert(self, tmp_path):
         assert result.exit_code == 0
         # Arbitrary check
         assert "variant_position" in result.stdout
+
+
+class TestPlinkEndToEnd:
+    def test_convert(self, tmp_path):
+        ts_path = "tests/data/plink/example"
+        zarr_path = tmp_path / "zarr"
+        runner = ct.CliRunner()
+        result = runner.invoke(
+            cli.plink2zarr,
+            f"convert {ts_path} {zarr_path}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        result = runner.invoke(
+            cli.vcf2zarr_main, f"inspect {zarr_path}", catch_exceptions=False
+        )
+        assert result.exit_code == 0
+        # Arbitrary check
+        assert "variant_position" in result.stdout
+
+
+@pytest.mark.parametrize(
+    "cmd",
+    [
+        main.bio2zarr,
+        cli.vcf2zarr_main,
+        cli.plink2zarr,
+        cli.vcfpartition,
+        cli.tskit2zarr,
+    ],
+)
+def test_version(cmd):
+    runner = ct.CliRunner()
+    result = runner.invoke(cmd, ["--version"], catch_exceptions=False)
+    s = f"version {provenance.__version__}\n"
+    assert result.stdout.endswith(s)

tests/test_plink.py

@@ -7,7 +7,7 @@
 import sgkit as sg
 import xarray.testing as xt
 
-from bio2zarr import plink
+from bio2zarr import plink, vcf
 
 
 def test_missing_dependency():
@@ -396,3 +396,24 @@ def test_variant_length(self, ds):
             ds.variant_length,
             [1, 1, 1, 1, 1, 1],
         )
+
+
+@pytest.mark.parametrize(
+    "prefix",
+    [
+        "tests/data/plink/example",
+    ],
+)
+def test_against_plinks_vcf_output(prefix, tmp_path):
+    vcf_path = prefix + ".vcf"
+    plink_zarr = tmp_path / "plink.zarr"
+    vcf_zarr = tmp_path / "vcf.zarr"
+    plink.convert(prefix, plink_zarr)
+    vcf.convert([vcf_path], vcf_zarr)
+    ds1 = sg.load_dataset(plink_zarr)
+    ds2 = (
+        sg.load_dataset(vcf_zarr)
+        .drop_dims("filters")
+        .drop_vars(["variant_quality", "variant_PR", "contig_length"])
+    )
+    xt.assert_equal(ds1, ds2)