Initial version of plink1.9-vcf command

Author: jeromekelleher

.github/workflows/ci.yml

@@ -52,7 +52,7 @@ jobs:
           channels: conda-forge,bioconda
       - name: Install dependencies
         run: |
-          conda install bcftools
+          conda install bcftools plink
           python -m pip install --upgrade pip
           python -m pip install '.[dev]'
           # Build the extension module in-place so pytest can find it
@@ -158,7 +158,7 @@ jobs:
           channels: conda-forge,bioconda
       - name: Install dependencies
         run: |
-          conda install bcftools
+          conda install bcftools plink
           python -m pip install --upgrade pip
           python -m pip install '.[dev]'
           # Build the extension module in-place so pytest can find it
@@ -189,7 +189,7 @@ jobs:
           channels: conda-forge,bioconda
       - name: Install dependencies
         run: |
-          conda install bcftools
+          conda install bcftools plink
           python -m pip install --upgrade pip
           python -m pip install '.[dev]'
           # Build the extension module in-place so pytest can find it

lib/tests.c

@@ -854,12 +854,153 @@ test_ftoa(void)
 
     for (j = 0; j < sizeof(cases) / sizeof(*cases); j++) {
         len = vcz_ftoa(buf, cases[j].val);
-        /* printf("j = %d %f->%s=='%s'\n", (int) j, cases[j].val, cases[j].expected, buf); */
+        /* printf("j = %d %f->%s=='%s'\n", (int) j, cases[j].val, cases[j].expected,
+         * buf); */
         CU_ASSERT_EQUAL_FATAL(len, strlen(cases[j].expected));
         CU_ASSERT_STRING_EQUAL_FATAL(buf, cases[j].expected);
     }
 }
 
+static void
+test_encode_plink_single_genotype(void)
+{
+    struct test_case {
+        int8_t genotype[2];
+        char expected;
+    };
+    // clang-format off
+    struct test_case cases[] = {
+        {{-1, -1}, VCZ_PLINK_MISSING},
+        {{-2, -1}, VCZ_PLINK_MISSING},
+        {{-1, -2}, VCZ_PLINK_MISSING},
+        {{-2, -2}, VCZ_PLINK_MISSING},
+        /* Unknown alleles are treated as missing */
+        {{2, 2}, VCZ_PLINK_MISSING},
+        {{-1, 2}, VCZ_PLINK_MISSING},
+        {{2, -1}, VCZ_PLINK_MISSING},
+        /* Half-calls are homozygous */
+        {{0, -2}, VCZ_PLINK_HOM_A2},
+        {{1, -2}, VCZ_PLINK_HOM_A1},
+        /* Nominal cases */
+        {{1, 0}, VCZ_PLINK_HET},
+        {{0, 1}, VCZ_PLINK_HET},
+        {{0, 0}, VCZ_PLINK_HOM_A2},
+        {{1, 1}, VCZ_PLINK_HOM_A1},
+    };
+    // clang-format on
+    size_t j;
+    char buf;
+    int8_t a12[] = { 1, 0 };
+
+    for (j = 0; j < sizeof(cases) / sizeof(*cases); j++) {
+        vcz_encode_plink(1, 1, cases[j].genotype, a12, &buf);
+        /* printf("j = %d [%d,%d]->%d==%d'\n", (int) j, cases[j].genotype[0], */
+        /*     cases[j].genotype[1], cases[j].expected, buf); */
+        CU_ASSERT_EQUAL_FATAL(buf, cases[j].expected);
+    }
+}
+
+static void
+test_encode_plink_example(void)
+{
+
+    const size_t num_variants = 3;
+    const size_t num_samples = 3;
+    int j;
+    // clang-format off
+    int8_t genotypes[] = {
+        0, 0, 0, 1, 0, 0,
+        1, 0, 1, 1, 0, -2,
+        1, 1, 0, 1, -1,-1,
+    };
+    // clang-format on
+    int8_t a12[] = { 1, 0, 1, 0, 1, 0 };
+    int8_t expected[] = { 59, 50, 24 };
+    char buf[3];
+
+    vcz_encode_plink(num_variants, num_samples, genotypes, a12, buf);
+    for (j = 0; j < 3; j++) {
+        /* printf("%d\n", buf[j]); */
+        CU_ASSERT_EQUAL_FATAL(buf[j], expected[j]);
+    }
+}
+
+static void
+test_encode_plink_single_genotype_vary_a12(void)
+{
+    struct test_case {
+        int8_t genotype[2];
+        int8_t a12[2];
+        char expected;
+    };
+    // clang-format off
+    struct test_case cases[] = {
+        {{-1, -1}, {0, 1}, VCZ_PLINK_MISSING},
+        {{0, -2}, {0, 1}, VCZ_PLINK_HOM_A1},
+        {{1, -2}, {0, 1}, VCZ_PLINK_HOM_A2},
+        {{0, 0}, {1, 2}, VCZ_PLINK_MISSING},
+        /* Allele 1 is missing */
+        {{0, 0}, {-1, 0}, VCZ_PLINK_HOM_A2},
+        {{1, 0}, {-1, 0}, VCZ_PLINK_MISSING},
+        {{0, 1}, {-1, 0}, VCZ_PLINK_MISSING},
+        {{1, 1}, {-1, 0}, VCZ_PLINK_MISSING},
+        {{0, -1}, {-1, 0}, VCZ_PLINK_MISSING},
+        {{-1, -1}, {-1, 0}, VCZ_PLINK_MISSING},
+        /* Nominal cases */
+        {{2, 3}, {2, 3}, VCZ_PLINK_HET},
+        {{3, 2}, {2, 3}, VCZ_PLINK_HET},
+        {{3, 3}, {2, 3}, VCZ_PLINK_HOM_A2},
+        {{2, 2}, {2, 3}, VCZ_PLINK_HOM_A1},
+    };
+    // clang-format on
+    size_t j;
+    char buf;
+
+    for (j = 0; j < sizeof(cases) / sizeof(*cases); j++) {
+        vcz_encode_plink(1, 1, cases[j].genotype, cases[j].a12, &buf);
+        /* printf("j = %d [%d,%d]->%d==%d'\n", (int) j, cases[j].genotype[0], */
+        /*     cases[j].genotype[1], cases[j].expected, buf); */
+        CU_ASSERT_EQUAL_FATAL(buf, cases[j].expected);
+    }
+}
+
+static void
+test_encode_plink_all_zeros_instance(size_t num_variants, size_t num_samples)
+{
+    int8_t *genotypes = calloc(num_variants * num_samples, 2);
+    int8_t *a12 = calloc(num_variants, 2);
+    char *buf = malloc(num_variants * num_samples / 4);
+    size_t j;
+
+    CU_ASSERT_FATAL(num_samples % 4 == 0);
+    CU_ASSERT_FATAL(genotypes != NULL);
+    CU_ASSERT_FATAL(a12 != NULL);
+    CU_ASSERT_FATAL(buf != NULL);
+
+    for (j = 0; j < num_variants; j++) {
+        a12[2 * j] = 1;
+    }
+    /* printf("variants=%d samples=%d\n", (int) num_variants, (int) num_samples); */
+    vcz_encode_plink(num_variants, num_samples, genotypes, a12, buf);
+    for (j = 0; j < num_variants * num_samples / 4; j++) {
+        /* printf("%d %d\n", (int) j, buf[j]); */
+        CU_ASSERT_EQUAL_FATAL(buf[j], -1);
+    }
+
+    free(genotypes);
+    free(a12);
+    free(buf);
+}
+
+static void
+test_encode_plink_all_zeros(void)
+{
+    test_encode_plink_all_zeros_instance(1, 4);
+    test_encode_plink_all_zeros_instance(10, 4);
+    test_encode_plink_all_zeros_instance(1, 400);
+    test_encode_plink_all_zeros_instance(100, 400);
+}
+
 /*=================================================
   Test suite management
   =================================================
@@ -972,6 +1113,11 @@ main(int argc, char **argv)
         { "test_itoa_pow10", test_itoa_pow10 },
         { "test_itoa_boundary", test_itoa_boundary },
         { "test_ftoa", test_ftoa },
+        { "test_encode_plink_single_genotype", test_encode_plink_single_genotype },
+        { "test_encode_plink_single_genotype_vary_a12",
+            test_encode_plink_single_genotype_vary_a12 },
+        { "test_encode_plink_example", test_encode_plink_example },
+        { "test_encode_plink_all_zeros", test_encode_plink_all_zeros },
         { NULL, NULL },
     };
     return test_main(tests, argc, argv);

lib/vcf_encoder.c

@@ -62,7 +62,6 @@ vcz_itoa(char *restrict buf, int64_t value)
     return p;
 }
 
-
 int
 vcz_ftoa(char *restrict buf, float value)
 {
@@ -998,3 +997,62 @@ vcz_variant_encoder_free(vcz_variant_encoder_t *self)
         self->format_fields = NULL;
     }
 }
+
+int
+vcz_encode_plink(size_t num_variants, size_t num_samples, const int8_t *genotypes,
+    const int8_t *a12_allele, char *buf)
+{
+    size_t j, k, variant_offset, byte_offset, bit_pos;
+    int8_t a, b, allele_1, allele_2, code;
+    int mask;
+    const size_t bytes_per_variant = (num_samples + 3) / 4;
+
+    variant_offset = 0;
+
+    memset(buf, 0, bytes_per_variant * num_variants);
+
+    for (j = 0; j < num_variants; j++) {
+        allele_1 = a12_allele[j * 2];
+        allele_2 = a12_allele[j * 2 + 1];
+        for (k = 0; k < num_samples; k++) {
+            code = VCZ_PLINK_MISSING;
+            a = genotypes[j * num_samples * 2 + k * 2];
+            b = genotypes[j * num_samples * 2 + k * 2 + 1];
+            if (b == -2) {
+                /* Treated as a haploid call by plink */
+                if (a == allele_1) {
+                    code = VCZ_PLINK_HOM_A1;
+                } else if (a == allele_2) {
+                    code = VCZ_PLINK_HOM_A2;
+                }
+            } else {
+                if (a == allele_1) {
+                    if (b == allele_1) {
+                        code = VCZ_PLINK_HOM_A1;
+                    } else if (b == allele_2) {
+                        code = VCZ_PLINK_HET;
+                    }
+                } else if (a == allele_2) {
+                    if (b == allele_2) {
+                        code = VCZ_PLINK_HOM_A2;
+                    } else if (b == allele_1) {
+                        code = VCZ_PLINK_HET;
+                    }
+                }
+                if ((allele_1 == -1)
+                    && (code == VCZ_PLINK_HOM_A1 || code == VCZ_PLINK_HET)) {
+                    code = VCZ_PLINK_MISSING;
+                }
+            }
+
+            /* printf("a=%d b=%d a1=%d a2=%d code = %d\n", a, b, allele_1, allele_2,
+             * code); */
+            byte_offset = variant_offset + k / 4;
+            bit_pos = (k % 4) * 2;
+            mask = ~(0x3 << bit_pos);
+            buf[byte_offset] = (char) (buf[byte_offset] & mask) | (code << bit_pos);
+        }
+        variant_offset += bytes_per_variant;
+    }
+    return 0;
+}

lib/vcf_encoder.h

@@ -109,3 +109,11 @@ int64_t vcz_variant_encoder_encode(
 
 int vcz_itoa(char *buf, int64_t v);
 int vcz_ftoa(char *buf, float v);
+
+
+#define VCZ_PLINK_HOM_A1 0x0  /* 00 */
+#define VCZ_PLINK_HOM_A2 0x3  /* 11 */
+#define VCZ_PLINK_HET 0x2     /* 10 */
+#define VCZ_PLINK_MISSING 0x1 /* 01 */
+int vcz_encode_plink(size_t num_variants, size_t num_samples, const int8_t *genotypes,
+    const int8_t *a12_allele, char *buf);

pyproject.toml

@@ -76,7 +76,7 @@ indent-width = 4
 [tool.ruff.lint]
 select = ["E", "F", "B", "W", "I", "N", "UP", "A", "RUF", "PT"]
 #Allow uppercase names for e.g. call_AD
-ignore = ["N806", "N802", "A001", "A002"]
+ignore = ["N806", "N802", "N803", "A001", "A002"]
 
 fixable = ["ALL"]
 unfixable = []

tests/test_cpython_interface.py

@@ -453,3 +453,83 @@ def test_attrs(self, name):
         encoder = cls.__new__(cls)
         with pytest.raises(SystemError, match="not initialised"):
             getattr(encoder, name)
+
+
+class TestEncodePlink:
+    def test_bad_num_arguments(self):
+        with pytest.raises(TypeError):
+            _vcztools.encode_plink()
+        with pytest.raises(TypeError):
+            _vcztools.encode_plink(0)
+
+    @pytest.mark.parametrize("bad_type", [[], {}, "string", 4])
+    def test_bad_types(self, bad_type):
+        array = np.zeros(5, dtype=np.int8)
+        with pytest.raises(TypeError):
+            _vcztools.encode_plink(bad_type, array)
+        with pytest.raises(TypeError):
+            _vcztools.encode_plink(array, bad_type)
+
+    def test_wrong_genotype_dims(self):
+        with pytest.raises(ValueError, match="genotypes has wrong dimension"):
+            _vcztools.encode_plink(
+                np.zeros((1, 1), dtype=np.int8), np.zeros((1, 2), dtype=np.int8)
+            )
+
+    def test_wrong_a12_allele_dims(self):
+        with pytest.raises(ValueError, match="a12_allele has wrong dimension"):
+            _vcztools.encode_plink(
+                np.zeros((1, 1, 2), dtype=np.int8), np.zeros((1, 2, 3), dtype=np.int8)
+            )
+
+    @pytest.mark.parametrize("bad_dtype", [np.int16, np.int64, np.float64, "S1"])
+    def test_bad_genotype_dtype(self, bad_dtype):
+        with pytest.raises(ValueError, match="Wrong dtype for genotypes"):
+            _vcztools.encode_plink(
+                np.zeros((1, 1, 2), dtype=bad_dtype), np.zeros((1, 2), dtype=np.int8)
+            )
+
+    @pytest.mark.parametrize("bad_dtype", [np.int16, np.int64, np.float64, "S1"])
+    def test_bad_a12_allele_dtype(self, bad_dtype):
+        with pytest.raises(ValueError, match="Wrong dtype for a12_allele"):
+            _vcztools.encode_plink(
+                np.zeros((1, 1, 2), dtype=np.int8), np.zeros((1, 2), dtype=bad_dtype)
+            )
+
+    @pytest.mark.parametrize("bad_ploidy", [1, 3])
+    def test_bad_ploidy(self, bad_ploidy):
+        with pytest.raises(ValueError, match="Only diploid genotypes"):
+            _vcztools.encode_plink(
+                np.zeros((1, 1, bad_ploidy), dtype=np.int8),
+                np.zeros((1, 2), dtype=np.int8),
+            )
+
+    @pytest.mark.parametrize("num_variants", [1, 3])
+    def test_num_variants_mismatch(self, num_variants):
+        with pytest.raises(ValueError, match="same first dimension"):
+            _vcztools.encode_plink(
+                np.zeros((2, 1, 2), dtype=np.int8),
+                np.zeros((num_variants, 2), dtype=np.int8),
+            )
+
+    @pytest.mark.parametrize("num_cols", [1, 3])
+    def test_bad_a12_allele_cols(self, num_cols):
+        with pytest.raises(ValueError, match="exactly 2 columns"):
+            _vcztools.encode_plink(
+                np.zeros((2, 1, 2), dtype=np.int8),
+                np.zeros((2, num_cols), dtype=np.int8),
+            )
+
+    def test_example(self):
+        G = np.array(
+            [
+                [[0, 0], [0, 1], [0, 0]],
+                [[1, 0], [1, 1], [0, -2]],
+                [[1, 1], [0, 1], [-1, -1]],
+            ],
+            dtype=np.int8,
+        )
+        a12_allele = np.array([[1, 0], [1, 0], [1, 0]], dtype=np.int8)
+        enc = _vcztools.encode_plink(G, a12_allele)
+        assert enc.dtype == np.uint8
+        assert list(enc) == [59, 50, 24]