Test on generated bed files written by bed_reader

Author: jeromekelleher

bio2zarr/plink.py

@@ -62,6 +62,8 @@ def __init__(self, path, num_variants, num_samples):
         # bytes per variant: 1 byte per 4 samples, rounded up
         self.bytes_per_variant = (self.num_samples + 3) // 4
 
+        # TODO open this as a persistent file and support reading from a
+        # stream
         with open(self.path, "rb") as f:
             magic = f.read(3)
             if magic != b"\x6c\x1b\x01":
@@ -132,10 +134,8 @@ def __init__(self, prefix):
             self.prefix + ".bim",
             self.prefix + ".fam",
         )
-
         self.bim = read_bim(self.paths.bim_path)
         self.fam = read_fam(self.paths.fam_path)
-
         self._num_records = self.bim.shape[0]
         self._num_samples = self.fam.shape[0]
         self.bed_reader = BedReader(
@@ -177,11 +177,8 @@ def iter_id(self, start, stop):
     def iter_alleles_and_genotypes(self, start, stop, shape, num_alleles):
         alt_field = self.bim.allele_1.values
         ref_field = self.bim.allele_2.values
-
         gt = self.bed_reader.decode(start, stop)
-
         phased = np.zeros(gt.shape[:2], dtype=bool)
-
         for i, (ref, alt) in enumerate(
             zip(ref_field[start:stop], alt_field[start:stop])
         ):

tests/test_plink.py

@@ -1,3 +1,5 @@
+import os.path
+
 import bed_reader
 import numpy as np
 import numpy.testing as nt
@@ -61,6 +63,51 @@ def test_bad_file_type(self, path):
         with pytest.raises(ValueError, match="Invalid BED file magic bytes"):
             plink.BedReader(path, 1, 1)
 
+    @pytest.mark.parametrize(
+        ("num_variants", "num_samples"),
+        [
+            (1, 1),
+            (1, 2),
+            (1, 3),
+            (1, 4),
+            (1, 5),
+            (1, 6),
+            (1, 7),
+            (1, 8),
+            (1, 9),
+            (2, 1),
+            (3, 1),
+            (10, 1),
+            (100, 1),
+            (10, 2),
+            (10, 3),
+            (10, 4),
+            (10, 5),
+            (20, 20),
+            (30, 3),
+        ],
+    )
+    def test_generated_bed_files(self, tmp_path, num_variants, num_samples):
+        bed_file = tmp_path / "a_file.bed"
+        # Generate a regular pattern of all possible values
+        data = np.arange(num_variants * num_samples, dtype=int) % 4
+        data[data == 3] = -127
+        data = data.reshape((num_variants, num_samples))
+
+        bed_reader.to_bed(bed_file, data.T, num_threads=1)
+
+        bytes_per_variant = (num_samples + 3) // 4
+        expected_size = 3 + bytes_per_variant * num_variants
+        assert os.path.getsize(bed_file) == expected_size
+
+        br_map = {0: (0, 0), 1: (0, 1), 2: (1, 1), -127: (-1, -1)}
+        reader = plink.BedReader(bed_file, num_variants, num_samples)
+        g = reader.decode(0, num_variants)
+        assert g.shape == (num_variants, num_samples, 2)
+        for j in range(num_variants):
+            for k in range(num_samples):
+                assert br_map[data[j, k]] == tuple(g[j, k])
+
 
 class TestSmallExample:
     @pytest.fixture(scope="class")