Fix plink data orientation error

jeromekelleher · jeromekelleher · commit d7122c94bbb8 · 2024-03-06T15:50:17.000Z
Plink data was not being correctly converted
diff --git a/bio2zarr/plink.py b/bio2zarr/plink.py
@@ -12,7 +12,11 @@
 
 
 def encode_genotypes_slice(bed_path, zarr_path, start, stop):
-    bed = bed_reader.open_bed(bed_path, num_threads=1)
+    # We need to count the A2 alleles here if we want to keep the
+    # alleles reported as allele_1, allele_2. It's obvious here what
+    # the correct approach is, but it is important to note that the
+    # 0th allele is *not* necessarily the REF for these datasets.
+    bed = bed_reader.open_bed(bed_path, num_threads=1, count_A1=False)
     store = zarr.DirectoryStore(zarr_path)
     root = zarr.group(store=store)
     gt = core.BufferedArray(root["call_genotype"], start)
@@ -62,7 +66,6 @@ def convert(
     bed = bed_reader.open_bed(bed_path, num_threads=1)
     n = bed.iid_count
     m = bed.sid_count
-    del bed
     logging.info(f"Scanned plink with {n} samples and {m} variants")
 
     # FIXME
@@ -79,6 +82,37 @@ def convert(
     chunks = [chunk_length, chunk_width]
     dimensions = ["variants", "samples"]
 
+    a = root.array(
+        "sample_id",
+        bed.iid,
+        dtype="str",
+        compressor=core.default_compressor,
+        chunks=(chunk_width,),
+    )
+    a.attrs["_ARRAY_DIMENSIONS"] = ["samples"]
+
+    # TODO encode these in slices - but read them in one go to avoid
+    # fetching repeatedly from bim file
+    a = root.array(
+        "variant_position",
+        bed.bp_position,
+        dtype=np.int32,
+        compressor=core.default_compressor,
+        chunks=(chunk_length,),
+    )
+    a.attrs["_ARRAY_DIMENSIONS"] = ["variants"]
+
+    alleles = np.stack([bed.allele_1, bed.allele_2], axis=1)
+    a = root.array(
+        "variant_allele",
+        alleles,
+        dtype="str",
+        compressor=core.default_compressor,
+        chunks=(chunk_length,),
+    )
+    a.attrs["_ARRAY_DIMENSIONS"] = ["variants", "alleles"]
+
+    # TODO remove this?
     a = root.empty(
         "call_genotype_phased",
         dtype="bool",
@@ -108,6 +142,8 @@ def convert(
     )
     a.attrs["_ARRAY_DIMENSIONS"] = list(dimensions)
 
+    del bed
+
     num_slices = max(1, worker_processes * 4)
     slices = core.chunk_aligned_slices(a, num_slices)
 
@@ -132,7 +168,7 @@ def validate(bed_path, zarr_path):
     root = zarr.group(store=store)
     call_genotype = root["call_genotype"][:]
 
-    bed = bed_reader.open_bed(bed_path, num_threads=1)
+    bed = bed_reader.open_bed(bed_path, count_A1=False, num_threads=1)
 
     assert call_genotype.shape[0] == bed.sid_count
     assert call_genotype.shape[1] == bed.iid_count
diff --git a/tests/test_plink.py b/tests/test_plink.py
@@ -3,22 +3,162 @@
 import xarray.testing as xt
 import pytest
 import sgkit as sg
+import sgkit.io.plink
+import bed_reader
 import zarr
 
 from bio2zarr import plink
 
 
 class TestSmallExample:
+    @pytest.fixture(scope="class")
+    def bed_path(self, tmp_path_factory):
+        tmp_path = tmp_path_factory.mktemp("data")
+        path = tmp_path / "example.bed"
+        # 7 sites x 3 samples
+        # These are counts of allele 1, so we have to flip them
+        # to get the values we expect
+        dosages = np.array(
+            [
+                [0, 1, 2],
+                [1, 0, 2],
+                [0, 0, 0],
+                [-127, 0, 0],
+                [2, 0, 0],
+                [-127, -127, -127],
+                [2, 2, 2],
+            ],
+            dtype=np.int8,
+        )
+        bed_reader.to_bed(path, dosages.T)
+        return path
+
+    @pytest.fixture(scope="class")
+    def ds(self, tmp_path_factory, bed_path):
+        tmp_path = tmp_path_factory.mktemp("data")
+        zarr_path = tmp_path / "example.plink.zarr"
+        plink.convert(bed_path, zarr_path)
+        return sg.load_dataset(zarr_path)
+
+    def test_genotypes(self, ds):
+        call_genotype = ds.call_genotype.values
+        assert call_genotype.shape == (7, 3, 2)
+        nt.assert_array_equal(
+            call_genotype,
+            [
+                [[1, 1], [1, 0], [0, 0]],
+                [[1, 0], [1, 1], [0, 0]],
+                [[1, 1], [1, 1], [1, 1]],
+                [[-1, -1], [1, 1], [1, 1]],
+                [[0, 0], [1, 1], [1, 1]],
+                [[-1, -1], [-1, -1], [-1, -1]],
+                [[0, 0], [0, 0], [0, 0]],
+            ],
+        )
+
+
+class TestEqualSgkit:
+    def test_simulated_example(self, tmp_path):
+        data_path = "tests/data/plink/"
+        bed_path = data_path + "plink_sim_10s_100v_10pmiss.bed"
+        fam_path = data_path + "plink_sim_10s_100v_10pmiss.fam"
+        bim_path = data_path + "plink_sim_10s_100v_10pmiss.bim"
+        # print(bed_path)
+        # print(fam_path)
+        sg_ds = sgkit.io.plink.read_plink(
+            bed_path=bed_path, fam_path=fam_path, bim_path=bim_path
+        )
+        out = tmp_path / "example.plink.zarr"
+        plink.convert(bed_path, out)
+        ds = sg.load_dataset(out)
+        nt.assert_array_equal(ds.call_genotype.values, sg_ds.call_genotype.values)
+
+
+class TestExample:
+    """
+    .bim file looks like this:
+
+    1       1_10    0       10      A       G
+    1       1_20    0       20      T       C
+
+    Definition: https://www.cog-genomics.org/plink/1.9/formats#bim
+    Chromosome code (either an integer, or 'X'/'Y'/'XY'/'MT'; '0'
+        indicates unknown) or name
+    Variant identifier
+    Position in morgans or centimorgans (safe to use dummy value of '0')
+    Base-pair coordinate (1-based; limited to 231-2)
+    Allele 1 (corresponding to clear bits in .bed; usually minor)
+    Allele 2 (corresponding to set bits in .bed; usually major)
+    """
+
+    @pytest.fixture(scope="class")
+    def ds(self, tmp_path_factory):
+        path = "tests/data/plink/example.bed"
+        out = tmp_path_factory.mktemp("data") / "example.plink.zarr"
+        plink.convert(path, out)
+        return sg.load_dataset(out)
+
+    def test_sample_ids(self, ds):
+        nt.assert_array_equal(ds.sample_id, [f"ind{j}" for j in range(10)])
+
+    def test_variant_position(self, ds):
+        nt.assert_array_equal(ds.variant_position, [10, 20])
+
+    def test_variant_allele(self, ds):
+        nt.assert_array_equal(ds.variant_allele, [["A", "G"], ["T", "C"]])
+
+    def test_genotypes(self, ds):
+        call_genotype = ds.call_genotype.values
+        assert call_genotype.shape == (2, 10, 2)
+        nt.assert_array_equal(
+            call_genotype,
+            [
+                [
+                    [0, 0],
+                    [0, 0],
+                    [0, 0],
+                    [1, 1],
+                    [1, 1],
+                    [1, 1],
+                    [1, 1],
+                    [1, 1],
+                    [1, 1],
+                    [1, 1],
+                ],
+                [
+                    [0, 0],
+                    [0, 0],
+                    [0, 0],
+                    [0, 0],
+                    [1, 1],
+                    [1, 1],
+                    [1, 1],
+                    [1, 1],
+                    [1, 1],
+                    [1, 1],
+                ],
+            ],
+        )
+
+    def test_sgkit(self, ds):
+        path = "tests/data/plink/example"
+        sg_ds = sgkit.io.plink.read_plink(path=path)
+        # Can't compare the full dataset yet
+        nt.assert_array_equal(ds.call_genotype.values, sg_ds.call_genotype.values)
+        # https://github.com/pystatgen/sgkit/issues/1209
+        nt.assert_array_equal(ds.variant_allele, sg_ds.variant_allele.astype(str))
+        nt.assert_array_equal(ds.variant_position, sg_ds.variant_position)
+        nt.assert_array_equal(ds.sample_id, sg_ds.sample_id)
+
+
+class TestSimulatedExample:
     @pytest.fixture(scope="class")
     def ds(self, tmp_path_factory):
         path = "tests/data/plink/plink_sim_10s_100v_10pmiss.bed"
         out = tmp_path_factory.mktemp("data") / "example.plink.zarr"
         plink.convert(path, out)
         return sg.load_dataset(out)
 
-    @pytest.mark.xfail
-    # FIXME I'm not sure these are the correct genotypes here, at least
-    # the test isn't passing and others are
     def test_genotypes(self, ds):
         # Validate a few randomly selected individual calls
         # (spanning all possible states for a call)
@@ -82,13 +222,6 @@ def test_chunk_size(
             worker_processes=worker_processes,
         )
         ds2 = sg.load_dataset(out)
-        # print()
-        # print(ds.call_genotype.values[2])
-        # print(ds2.call_genotype.values[2])
-
-        # print(ds2)
-        # print(ds2.call_genotype.values)
-        # print(ds.call_genotype.values)
         xt.assert_equal(ds, ds2)
         # TODO check array chunks
 
@@ -101,8 +234,6 @@ def test_chunk_size(
         (33, 3),
         (99, 10),
         (3, 10),
-        # This one doesn't fail as it's the same as defaults
-        # (100, 10),
     ],
 )
 @pytest.mark.parametrize("worker_processes", [0])
diff --git a/tests/test_simulated_data.py b/tests/test_simulated_data.py
@@ -2,6 +2,7 @@
 import msprime
 import pysam
 import sgkit as sg
+import numpy as np
 import numpy.testing as nt
 
 from bio2zarr import vcf
@@ -41,7 +42,3 @@ def test_ploidy(self, ploidy, tmp_path):
         nt.assert_equal(ds.variant_allele[:, 0].values, "A")
         nt.assert_equal(ds.variant_allele[:, 1].values, "T")
         nt.assert_equal(ds.variant_position, ts.sites_position)
-
-
-# TODO add a plink equivalant using
-# https://fastlmm.github.io/bed-reader/#bed_reader.to_bed
