Show forward and backward probability matrices from tskit and BEAGLE

tskit-dev · Aug 9, 2023 · dc0d3e9 · dc0d3e9
1 parent cb35bb3
commit dc0d3e9
Showing 1 changed file with 13 additions and 6 deletions.
diff --git a/python/tests/test_imputation.py b/python/tests/test_imputation.py
@@ -286,9 +286,16 @@ def test_toy_example():
     print(list(ref_ts.haplotypes()))
     print(ref_ts)
     print(query)
-    fw, bw = get_tskit_forward_backward_matrices(ref_ts, query[0])
-    print(fw)
-    print(bw)
-    df_beagle = parse_matrix(beagle_fwd_matrix_text_1)
-    # compare beagle results with tskit
-    print(df_beagle["val"])
+    tskit_fwd, tskit_bwd = get_tskit_forward_backward_matrices(ref_ts, query[0])
+    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)))