Import BEAGLE implementation

python/tests/test_imputation.py

@@ -7,6 +7,7 @@
 
 import _tskit
 import tskit
+from tests.beagle import run_beagle
 
 
 # A toy tree sequence containing 3 diploid individuals with 5 sites and 5 mutations.
@@ -289,20 +290,49 @@ def parse_matrix(csv_text):
 
 
 def test_toy_example():
+    # Get toy dataset ready for tskit.
     ref_ts, query = get_toy_data()
     print(list(ref_ts.haplotypes()))
     print(ref_ts)
     print(query)
-    tskit_fwd, tskit_bwd = get_tskit_forward_backward_matrices(ref_ts, query[0])
+    # tskit_fwd, tskit_bwd = get_tskit_forward_backward_matrices(ref_ts, query[0])
+    # Get BEAGLE results for the toy dataset.
     beagle_fwd = parse_matrix(beagle_fwd_matrix_text_1)
     beagle_bwd = parse_matrix(beagle_bwd_matrix_text_1)
-    print("Forward probability matrix")
-    print("tskit")
-    print(tskit_fwd)
-    print("beagle")
-    print(beagle_fwd["val"].reshape((4, 6)))
-    print("Backward probability matrix")
-    print("tskit")
-    print(tskit_bwd)
-    print("beagle")
-    print(beagle_bwd["val"].reshape((4, 6)))
+    print(beagle_fwd)
+    print(beagle_bwd)
+    # Run a trivial test case using Python implementation of BEAGLE.
+    ref_h = np.array(
+        [
+            [
+                0,
+                1,
+                0,
+                1,
+                0,
+                1,
+            ],
+            [
+                1,
+                0,
+                1,
+                0,
+                1,
+                0,
+            ],
+        ]
+    ).T
+    query_h = np.array(
+        [
+            -1,
+            1,
+            -1,
+            1,
+            -1,
+            0,
+        ]
+    )
+    assert ref_h.shape[0] == len(query_h)
+    pos = (np.arange(len(query_h)) + 1) * 1e4
+    ia = run_beagle(ref_h, query_h, pos, miscall_rate=0.1, ne=10.0, debug=True)
+    print(ia)