diff --git a/python/tests/test_genotype_matching.py b/python/tests/test_genotype_matching.py index a04e3873a6..ce35b90b18 100644 --- a/python/tests/test_genotype_matching.py +++ b/python/tests/test_genotype_matching.py @@ -1345,12 +1345,17 @@ def verify(self, ts): self.assertAllClose(ll_tree, ll_mirror_tree_dict) # Ensure that the decoded matrices are the same + flipped_G_check = np.flip(G_check, axis=0) + flipped_s = np.flip(s, axis=1) + num_alleles = ls.core.get_num_alleles(flipped_G_check, flipped_s) + F_mirror_matrix, c, ll = ls.forwards( - np.flip(G_check, axis=0), - np.flip(s, axis=1), - r_flip, - p_mutation=np.flip(mu), - scale_mutation_based_on_n_alleles=False, + flipped_G_check, + flipped_s, + num_alleles=num_alleles, + prob_recombination=r_flip, + prob_mutation=np.flip(mu), + scale_mutation_rate=False, ) self.assertAllClose(F_mirror_matrix, cm_mirror.decode()) @@ -1373,8 +1378,14 @@ def verify(self, ts): ts_check.genotype_matrix()[i, :], ts_check.genotype_matrix()[i, :] ) + num_alleles = ls.core.get_num_alleles(G_check, s) F, c, ll = ls.forwards( - G_check, s, r, p_mutation=mu, scale_mutation_based_on_n_alleles=False + reference_panel=G_check, + query=s, + num_alleles=num_alleles, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=False, ) cm_d = ls_forward_tree(s[0, :], ts_check, r, mu) self.assertAllClose(cm_d.decode(), F) @@ -1399,16 +1410,23 @@ def verify(self, ts): ts_check.genotype_matrix()[i, :], ts_check.genotype_matrix()[i, :] ) + num_alleles = ls.core.get_num_alleles(G_check, s) F, c, ll = ls.forwards( - G_check, s, r, p_mutation=mu, scale_mutation_based_on_n_alleles=False + reference_panel=G_check, + query=s, + num_alleles=num_alleles, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=False, ) B = ls.backwards( - G_check, - s, - c, - r, - p_mutation=mu, - scale_mutation_based_on_n_alleles=False, + reference_panel=G_check, + query=s, + num_alleles=num_alleles, + normalisation_factor_from_forward=c, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=False, ) # Note, need to remove the first sample from the ts, and ensure that @@ -1453,16 +1471,24 @@ def verify(self, ts): ts_check.genotype_matrix()[i, :], ts_check.genotype_matrix()[i, :] ) ts_check = ts.simplify(range(1, n + 1), filter_sites=False) + + num_alleles = ls.core.get_num_alleles(G_check, s) phased_path, ll = ls.viterbi( - G_check, s, r, p_mutation=mu, scale_mutation_based_on_n_alleles=False + reference_panel=G_check, + query=s, + num_alleles=num_alleles, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=False, ) - path_ll_matrix = ls.path_ll( - G_check, - s, - phased_path, - r, - p_mutation=mu, - scale_mutation_based_on_n_alleles=False, + path_ll_matrix = ls.path_loglik( + reference_panel=G_check, + query=s, + num_alleles=num_alleles, + path=phased_path, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=False, ) c_v = ls_viterbi_tree(s[0, :], ts_check, r, mu) @@ -1471,13 +1497,14 @@ def verify(self, ts): # Attempt to get the path path_tree_dict = c_v.traceback() # Work out the likelihood of the proposed path - path_ll_tree = ls.path_ll( - G_check, - s, - np.transpose(path_tree_dict), - r, - p_mutation=mu, - scale_mutation_based_on_n_alleles=False, + path_ll_tree = ls.path_loglik( + reference_panel=G_check, + query=s, + num_alleles=num_alleles, + path=np.transpose(path_tree_dict), + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=False, ) self.assertAllClose(ll, ll_tree) diff --git a/python/tests/test_haplotype_matching.py b/python/tests/test_haplotype_matching.py index dcc1d684fb..6efb41d93b 100644 --- a/python/tests/test_haplotype_matching.py +++ b/python/tests/test_haplotype_matching.py @@ -925,12 +925,16 @@ def verify(self, ts): self.assertAllClose(ll_tree, ll_mirror_tree) # Ensure that the decoded matrices are the same + flipped_H = np.flip(H, axis=0) + flipped_s = np.flip(s, axis=1) + num_alleles = ls.core.get_num_alleles(flipped_H, flipped_s) F_mirror_matrix, c, ll = ls.forwards( - np.flip(H, axis=0), - np.flip(s, axis=1), - r_flip, - p_mutation=np.flip(mu), - scale_mutation_based_on_n_alleles=False, + reference_panel=flipped_H, + query=flipped_s, + num_alleles=num_alleles, + prob_recombination=r_flip, + prob_mutation=np.flip(mu), + scale_mutation_rate=False, ) self.assertAllClose(F_mirror_matrix, cm_mirror.decode()) @@ -949,12 +953,14 @@ def verify(self, ts): # Warning from lshmm: # Passed a vector of mutation rates, but rescaling each mutation # rate conditional on the number of alleles + num_alleles = ls.core.get_num_alleles(H, s) F, c, ll = ls.forwards( - H, - s, - r, - p_mutation=mu, - scale_mutation_based_on_n_alleles=scale_mutation, + reference_panel=H, + query=s, + num_alleles=num_alleles, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=scale_mutation, ) # Note, need to remove the first sample from the ts, and ensure # that invariant sites aren't removed. @@ -977,16 +983,23 @@ class TestForwardBackwardTree(FBAlgorithmBase): def verify(self, ts): for n, H, s, r, mu in self.example_parameters_haplotypes(ts): + num_alleles = ls.core.get_num_alleles(H, s) F, c, ll = ls.forwards( - H, s, r, p_mutation=mu, scale_mutation_based_on_n_alleles=False + reference_panel=H, + query=s, + num_alleles=num_alleles, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=False, ) B = ls.backwards( - H, - s, - c, - r, - p_mutation=mu, - scale_mutation_based_on_n_alleles=False, + reference_panel=H, + query=s, + num_alleles=num_alleles, + normalisation_factor_from_forward=c, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=False, ) # Note, need to remove the first sample from the ts, and ensure that @@ -1017,8 +1030,14 @@ class TestTreeViterbiHap(VitAlgorithmBase): def verify(self, ts): for n, H, s, r, mu in self.example_parameters_haplotypes(ts): + num_alleles = ls.core.get_num_alleles(H, s) path, ll = ls.viterbi( - H, s, r, p_mutation=mu, scale_mutation_based_on_n_alleles=False + reference_panel=H, + query=s, + num_alleles=num_alleles, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=False, ) ts_check = ts.simplify(range(1, n + 1), filter_sites=False) cm = ls_viterbi_tree(s[0, :], ts_check, r, mu) @@ -1028,13 +1047,14 @@ def verify(self, ts): # Now, need to ensure that the likelihood of the preferred path is # the same as ll_tree (and ll). path_tree = cm.traceback() - ll_check = ls.path_ll( - H, - s, - path_tree, - r, - p_mutation=mu, - scale_mutation_based_on_n_alleles=False, + ll_check = ls.path_loglik( + reference_panel=H, + query=s, + num_alleles=num_alleles, + path=path_tree, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=False, ) self.assertAllClose(ll, ll_check) @@ -1051,13 +1071,16 @@ def check_viterbi(ts, h, recombination=None, mutation=None): precision = 22 G = ts.genotype_matrix() + s = h.reshape(1, m) + num_alleles = ls.core.get_num_alleles(G, s) path, ll = ls.viterbi( - G, - h.reshape(1, m), - recombination, - p_mutation=mutation, - scale_mutation_based_on_n_alleles=False, + reference_panel=G, + query=s, + num_alleles=num_alleles, + prob_recombination=recombination, + prob_mutation=mutation, + scale_mutation_rate=False, ) assert np.isscalar(ll) @@ -1069,13 +1092,14 @@ def check_viterbi(ts, h, recombination=None, mutation=None): # Check that the likelihood of the preferred path is # the same as ll_tree (and ll). path_tree = cm.traceback() - ll_check = ls.path_ll( - G, - h.reshape(1, m), - path_tree, - recombination, - p_mutation=mutation, - scale_mutation_based_on_n_alleles=False, + ll_check = ls.path_loglik( + reference_panel=G, + query=s, + num_allele=num_alleles, + path=path_tree, + prob_recombination=recombination, + prob_mutation=mutation, + scale_mutation_rate=False, ) nt.assert_allclose(ll_check, ll) @@ -1106,12 +1130,16 @@ def check_forward_matrix(ts, h, recombination=None, mutation=None): mutation = np.zeros(ts.num_sites) G = ts.genotype_matrix() + s = h.reshape(1, m) + num_alleles = ls.core.get_num_alleles(G, s) + F, c, ll = ls.forwards( - G, - h.reshape(1, m), - recombination, - p_mutation=mutation, - scale_mutation_based_on_n_alleles=False, + reference_panel=G, + query=s, + num_alleles=num_alleles, + prob_recombination=recombination, + prob_mutation=mutation, + scale_mutation_rate=False, ) assert F.shape == (m, n) assert c.shape == (m,) @@ -1150,13 +1178,17 @@ def check_backward_matrix(ts, h, forward_cm, recombination=None, mutation=None): mutation = np.zeros(ts.num_sites) G = ts.genotype_matrix() + s = h.reshape(1, m) + num_alleles = ls.core.get_num_alleles(G, s) + B = ls.backwards( - G, - h.reshape(1, m), - forward_cm.normalisation_factor, - recombination, - p_mutation=mutation, - scale_mutation_based_on_n_alleles=False, + reference_pane=G, + query=h.reshape(1, m), + num_alleles=num_alleles, + normalisation_factor_from_forward=forward_cm.normalisation_factor, + prob_recombination=recombination, + prob_mutation=mutation, + scale_mutation_rate=False, ) backward_cm = ls_backward_tree(