Modify toy examples 7 and 8 to include ACGT alleles

python/tests/test_imputation.py

@@ -20,8 +20,6 @@
 # In this setup, A is encoded as 0 and C as 1 whether the allele encoding is
 # ancestral/derived (0/1) or ACGT (0123).
 #
-# TODO: Add more test cases covering imputation of all the ACGT alleles.
-#
 # Case 0:
 # Reference haplotypes and query haplotypes have 0|0 at all sites.
 # First and last reference markers are missing in the query haplotypes.
@@ -374,28 +372,28 @@
 # fmt: off
 toy_ref_7 = np.array(
     [
-        [0, 0, 0, 1, 1, 0, 0, 0, 1],    # ref_0
-        [1, 0, 0, 0, 1, 0, 0, 0, 1],    # ref_0
-        [0, 0, 1, 1, 1, 0, 0, 0, 1],    # ref_1
-        [0, 1, 0, 0, 0, 0, 0, 1, 0],    # ref_1
-        [0, 0, 0, 0, 0, 1, 1, 1, 0],    # ref_2
-        [1, 0, 0, 0, 1, 0, 0, 0, 1],    # ref_2
-        [0, 0, 1, 1, 1, 0, 0, 0, 1],    # ref_3
-        [0, 0, 0, 0, 0, 0, 0, 0, 1],    # ref_3
-        [0, 1, 0, 0, 0, 0, 0, 1, 0],    # ref_4
-        [0, 0, 0, 0, 0, 1, 1, 1, 0],    # ref_4
+        [0, 1, 2, 0, 1, 1, 2, 3, 1],    # ref_0
+        [1, 1, 2, 3, 1, 1, 2, 3, 1],    # ref_0
+        [0, 1, 3, 0, 1, 1, 2, 3, 1],    # ref_1
+        [0, 2, 2, 3, 0, 1, 2, 0, 0],    # ref_1
+        [0, 1, 2, 3, 0, 2, 3, 0, 0],    # ref_2
+        [1, 1, 2, 3, 1, 1, 2, 3, 1],    # ref_2
+        [0, 1, 3, 0, 1, 1, 2, 3, 1],    # ref_3
+        [0, 1, 2, 3, 0, 1, 2, 3, 1],    # ref_3
+        [0, 2, 2, 3, 0, 1, 2, 0, 0],    # ref_4
+        [0, 1, 2, 3, 0, 2, 3, 0, 0],    # ref_4
     ], dtype=np.int32
 ).T
 toy_query_7 = np.array(
     [
-        [0, 0, -1, 1, -1, 1, -1, 0, 1],   # query_0, haplotype 0
-        [1, 1, -1, 0, -1, 0, -1, 1, 0],   # query_0, haplotype 1
+        [0, 1, -1, 0, -1, 2, -1, 3, 1],   # query_0, haplotype 0
+        [1, 2, -1, 3, -1, 1, -1, 0, 0],   # query_0, haplotype 1
     ], dtype=np.int32
 )
 toy_query_7_beagle_imputed = np.array(
     [
-        [0, 0, 1, 1, 1, 1, 0, 0, 1],    # query_0, haplotype 0
-        [1, 1, 0, 0, 0, 0, 0, 1, 0],    # query_0, haplotype 1
+        [0, 1, 3, 0, 1, 2, 2, 3, 1],    # query_0, haplotype 0
+        [1, 2, 2, 3, 0, 1, 2, 0, 0],    # query_0, haplotype 1
     ], dtype=np.int32
 )
 toy_query_7_beagle_allele_probs = np.array(
@@ -427,14 +425,14 @@
 toy_ref_8 = np.copy(toy_ref_7)
 toy_query_8 = np.array(
     [
-        [-1, 0, -1, 1, -1, 1, -1, 0, -1],  # query_0, haplotype 0
-        [-1, 1, -1, 0, -1, 0, -1, 1, -1],  # query_0, haplotype 1
+        [-1, 1, -1, 0, -1, 2, -1, 3, -1],  # query_0, haplotype 0
+        [-1, 2, -1, 3, -1, 1, -1, 0, -1],  # query_0, haplotype 1
     ], dtype=np.int32
 )
 toy_query_8_beagle_imputed = np.array(
     [
-        [0, 0, 1, 1, 1, 1, 0, 0, 1],    # query_0, haplotype 0
-        [0, 1, 0, 0, 0, 0, 0, 1, 0],    # query_0, haplotype 1
+        [0, 1, 3, 0, 1, 2, 2, 3, 1],    # query_0, haplotype 0
+        [0, 2, 2, 3, 0, 1, 2, 0, 0],    # query_0, haplotype 1
     ], dtype=np.int32
 )
 toy_query_8_beagle_allele_probs = np.array(
@@ -581,26 +579,26 @@ def test_beagle_numba(input_ref, input_query):
 toy_ts_sites_text = """\
 position        ancestral_state metadata
 10000.000000    A
-20000.000000    A
-30000.000000    A
-40000.000000    A
+20000.000000    C
+30000.000000    G
+40000.000000    T
 50000.000000    A
-60000.000000    A
-70000.000000    A
-80000.000000    A
+60000.000000    C
+70000.000000    G
+80000.000000    T
 90000.000000    A
 """
 
 toy_ts_mutations_text = """\
 site    node    time    derived_state   parent  metadata
 0       10      unknown C       -1
-1       11      unknown C       -1
-2       12      unknown C       -1
-3       13      unknown C       -1
+1       11      unknown G       -1
+2       12      unknown T       -1
+3       13      unknown A       -1
 4       14      unknown C       -1
-5       16      unknown C       -1
-6       16      unknown C       -1
-7       17      unknown C       -1
+5       16      unknown G       -1
+6       16      unknown T       -1
+7       17      unknown A       -1
 8       19      unknown C       -1
 """