From cbfc44d5212d9233c07c7aefe2b13cc834af8b55 Mon Sep 17 00:00:00 2001 From: Shing Zhan Date: Fri, 1 Sep 2023 20:28:49 +0100 Subject: [PATCH] Import BEAGLE implementation --- python/tests/test_imputation.py | 52 ++++++++++++++++++++++++++------- 1 file changed, 41 insertions(+), 11 deletions(-) diff --git a/python/tests/test_imputation.py b/python/tests/test_imputation.py index 1267ee92fb..d4ce26b189 100644 --- a/python/tests/test_imputation.py +++ b/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)