Fixup tests for plink

Covert some corner cases in the plink outout

Author: jeromekelleher

bio2zarr/plink.py

@@ -3,7 +3,6 @@
 import pathlib
 
 import numpy as np
-import zarr
 
 from bio2zarr import constants, core, vcz
 
@@ -216,41 +215,3 @@ def convert(
     )
     vzw.finalise(show_progress)
     vzw.create_index()
-
-
-# FIXME do this more efficiently - currently reading the whole thing
-# in for convenience, and also comparing call-by-call
-# TODO we should remove this function from the API - it's a test function
-# and should be moved into the suite
-@core.requires_optional_dependency("bed_reader", "plink")
-def validate(bed_path, zarr_path):
-    import bed_reader
-
-    root = zarr.open(store=zarr_path, mode="r")
-    call_genotype = root["call_genotype"][:]
-
-    bed = bed_reader.open_bed(bed_path + ".bed", count_A1=True, num_threads=1)
-
-    assert call_genotype.shape[0] == bed.sid_count
-    assert call_genotype.shape[1] == bed.iid_count
-    bed_genotypes = bed.read(dtype="int8").T
-    assert call_genotype.shape[0] == bed_genotypes.shape[0]
-    assert call_genotype.shape[1] == bed_genotypes.shape[1]
-    assert call_genotype.shape[2] == 2
-
-    row_id = 0
-    for bed_row, zarr_row in zip(bed_genotypes, call_genotype):
-        # print("ROW", row_id)
-        # print(bed_row, zarr_row)
-        row_id += 1
-        for bed_call, zarr_call in zip(bed_row, zarr_row):
-            if bed_call == -127:
-                assert list(zarr_call) == [-1, -1]
-            elif bed_call == 0:
-                assert list(zarr_call) == [0, 0]
-            elif bed_call == 1:
-                assert list(zarr_call) == [1, 0]
-            elif bed_call == 2:
-                assert list(zarr_call) == [1, 1]
-            else:  # pragma no cover
-                raise AssertionError(f"Unexpected bed call {bed_call}")

tests/test_plink.py

@@ -6,6 +6,7 @@
 import pytest
 import sgkit as sg
 import xarray.testing as xt
+import zarr
 
 from bio2zarr import plink, vcf
 
@@ -69,8 +70,8 @@ def test_genotypes(self, ds):
         nt.assert_array_equal(
             call_genotype,
             [
-                [[0, 0], [1, 0], [1, 1]],
-                [[1, 0], [0, 0], [1, 1]],
+                [[0, 0], [0, 1], [1, 1]],
+                [[0, 1], [0, 0], [1, 1]],
                 [[0, 0], [0, 0], [0, 0]],
                 [[-1, -1], [0, 0], [0, 0]],
                 [[1, 1], [0, 0], [0, 0]],
@@ -134,7 +135,7 @@ def test_variant_allele(self, ds):
         nt.assert_array_equal(ds.variant_allele, [["GG", "A"], ["C", "TTT"]])
 
     def test_variant_length(self, ds):
-        nt.assert_array_equal(ds.variant_length, [2, 3])
+        nt.assert_array_equal(ds.variant_length, [2, 1])
 
     def test_contig_id(self, ds):
         """Test that contig identifiers are correctly extracted and stored."""
@@ -207,8 +208,8 @@ def test_genotypes(self, ds):
         )
         expected = np.array(
             [
-                [1, 0],
-                [1, 0],
+                [0, 1],
+                [0, 1],
                 [1, 1],
                 [1, 1],
                 [-1, -1],
@@ -276,6 +277,37 @@ def test_chunk_size(
         # TODO check array chunks
 
 
+def validate(bed_path, zarr_path):
+    root = zarr.open(store=zarr_path, mode="r")
+    call_genotype = root["call_genotype"][:]
+
+    bed = bed_reader.open_bed(bed_path + ".bed", count_A1=True, num_threads=1)
+
+    assert call_genotype.shape[0] == bed.sid_count
+    assert call_genotype.shape[1] == bed.iid_count
+    bed_genotypes = bed.read(dtype="int8").T
+    assert call_genotype.shape[0] == bed_genotypes.shape[0]
+    assert call_genotype.shape[1] == bed_genotypes.shape[1]
+    assert call_genotype.shape[2] == 2
+
+    row_id = 0
+    for bed_row, zarr_row in zip(bed_genotypes, call_genotype):
+        # print("ROW", row_id)
+        # print(bed_row, zarr_row)
+        row_id += 1
+        for bed_call, zarr_call in zip(bed_row, zarr_row):
+            if bed_call == -127:
+                assert list(zarr_call) == [-1, -1]
+            elif bed_call == 0:
+                assert list(zarr_call) == [0, 0]
+            elif bed_call == 1:
+                assert list(zarr_call) == [0, 1]
+            elif bed_call == 2:
+                assert list(zarr_call) == [1, 1]
+            else:  # pragma no cover
+                raise AssertionError(f"Unexpected bed call {bed_call}")
+
+
 @pytest.mark.parametrize(
     ("variants_chunk_size", "samples_chunk_size"),
     [
@@ -299,7 +331,7 @@ def test_by_validating(
         samples_chunk_size=samples_chunk_size,
         worker_processes=worker_processes,
     )
-    plink.validate(path, out)
+    validate(path, out)
 
 
 class TestMultipleContigs:
@@ -379,11 +411,11 @@ def test_genotypes(self, ds):
         nt.assert_array_equal(
             call_genotype,
             [
-                [[0, 0], [1, 0], [1, 1], [0, 0]],  # chr1
-                [[1, 0], [0, 0], [1, 1], [1, 0]],  # chr1
+                [[0, 0], [0, 1], [1, 1], [0, 0]],  # chr1
+                [[0, 1], [0, 0], [1, 1], [0, 1]],  # chr1
                 [[0, 0], [0, 0], [0, 0], [1, 1]],  # chr2
-                [[1, 1], [0, 0], [1, 0], [0, 0]],  # chrX
-                [[1, 0], [1, 0], [1, 0], [1, 0]],  # chrX
+                [[1, 1], [0, 0], [0, 1], [0, 0]],  # chrX
+                [[0, 1], [0, 1], [0, 1], [0, 1]],  # chrX
                 [[0, 0], [1, 1], [0, 0], [1, 1]],  # chrY
             ],
         )