From 01736b2436651c09e4c2521c8452b6869c0c3537 Mon Sep 17 00:00:00 2001 From: szhan Date: Mon, 1 Jul 2024 18:18:30 +0100 Subject: [PATCH] Reference using argument names --- tests/test_nontree_fb_haploid.py | 42 ++++++++- tests/test_nontree_vit_haploid.py | 152 +++++++++++++++++++++++++----- 2 files changed, 168 insertions(+), 26 deletions(-) diff --git a/tests/test_nontree_fb_haploid.py b/tests/test_nontree_fb_haploid.py index c8f277e..331aabe 100644 --- a/tests/test_nontree_fb_haploid.py +++ b/tests/test_nontree_fb_haploid.py @@ -19,13 +19,47 @@ def verify(self, ts, scale_mutation_rate, include_ancestors): include_extreme_rates=True, ): emission_func = core.get_emission_probability_haploid - F_vs, c_vs, ll_vs = fbh.forwards_ls_hap(n, m, H_vs, s, e_vs, r, emission_func, norm=False) - B_vs = fbh.backwards_ls_hap(n, m, H_vs, s, e_vs, c_vs, r, emission_func) + F_vs, c_vs, ll_vs = fbh.forwards_ls_hap( + n=n, + m=m, + H=H_vs, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, + norm=False, + ) + B_vs = fbh.backwards_ls_hap( + n=n, + m=m, + H=H_vs, + s=s, + e=e_vs, + c=c_vs, + r=r, + emission_func=emission_func, + ) self.assertAllClose(np.log10(np.sum(F_vs * B_vs, 1)), ll_vs * np.ones(m)) F_tmp, c_tmp, ll_tmp = fbh.forwards_ls_hap( - n, m, H_vs, s, e_vs, r, emission_func, norm=True, + n=n, + m=m, + H=H_vs, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, + norm=True, + ) + B_tmp = fbh.backwards_ls_hap( + n=n, + m=m, + H=H_vs, + s=s, + e=e_vs, + c=c_tmp, + r=r, + emission_func=emission_func, ) - B_tmp = fbh.backwards_ls_hap(n, m, H_vs, s, e_vs, c_tmp, r, emission_func) self.assertAllClose(np.sum(F_tmp * B_tmp, 1), np.ones(m)) self.assertAllClose(ll_vs, ll_tmp) diff --git a/tests/test_nontree_vit_haploid.py b/tests/test_nontree_vit_haploid.py index 1b61887..475d121 100644 --- a/tests/test_nontree_vit_haploid.py +++ b/tests/test_nontree_vit_haploid.py @@ -18,47 +18,140 @@ def verify(self, ts, scale_mutation_rate, include_ancestors): include_ancestors=include_ancestors, include_extreme_rates=True, ): - V_vs, P_vs, ll_vs = vh.forwards_viterbi_hap_naive(n, m, H_vs, s, e_vs, r) - path_vs = vh.backwards_viterbi_hap(m, V_vs[m - 1, :], P_vs) - ll_check = vh.path_ll_hap(n, m, H_vs, path_vs, s, e_vs, r) + emission_func = core.get_emission_probability_haploid + + V_vs, P_vs, ll_vs = vh.forwards_viterbi_hap_naive( + n=n, + m=m, + H=H_vs, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, + ) + path_vs = vh.backwards_viterbi_hap( + m=m, + V_last=V_vs[m - 1, :], + P=P_vs, + ) + ll_check = vh.path_ll_hap( + n=n, + m=m, + H=H_vs, + path=path_vs, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, + ) self.assertAllClose(ll_vs, ll_check) V_tmp, P_tmp, ll_tmp = vh.forwards_viterbi_hap_naive_vec( - n, m, H_vs, s, e_vs, r + n=n, + m=m, + H=H_vs, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, + ) + path_tmp = vh.backwards_viterbi_hap( + m=m, + V_last=V_tmp[m - 1, :], + P=P_tmp, + ) + ll_check = vh.path_ll_hap( + n=n, + m=m, + H=H_vs, + path=path_tmp, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, ) - path_tmp = vh.backwards_viterbi_hap(m, V_tmp[m - 1, :], P_tmp) - ll_check = vh.path_ll_hap(n, m, H_vs, path_tmp, s, e_vs, r) self.assertAllClose(ll_tmp, ll_check) self.assertAllClose(ll_vs, ll_tmp) V_tmp, P_tmp, ll_tmp = vh.forwards_viterbi_hap_naive_low_mem( - n, m, H_vs, s, e_vs, r + n=n, + m=m, + H=H_vs, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, + ) + path_tmp = vh.backwards_viterbi_hap(m=m, V_last=V_tmp, P=P_tmp) + ll_check = vh.path_ll_hap( + n=n, + m=m, + H=H_vs, + path=path_tmp, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, ) - path_tmp = vh.backwards_viterbi_hap(m, V_tmp, P_tmp) - ll_check = vh.path_ll_hap(n, m, H_vs, path_tmp, s, e_vs, r) self.assertAllClose(ll_tmp, ll_check) self.assertAllClose(ll_vs, ll_tmp) V_tmp, P_tmp, ll_tmp = vh.forwards_viterbi_hap_naive_low_mem_rescaling( - n, m, H_vs, s, e_vs, r + n=n, + m=m, + H=H_vs, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, + ) + path_tmp = vh.backwards_viterbi_hap(m=m, V_last=V_tmp, P=P_tmp) + ll_check = vh.path_ll_hap( + n=n, + m=m, + H=H_vs, + path=path_tmp, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, ) - path_tmp = vh.backwards_viterbi_hap(m, V_tmp, P_tmp) - ll_check = vh.path_ll_hap(n, m, H_vs, path_tmp, s, e_vs, r) self.assertAllClose(ll_tmp, ll_check) self.assertAllClose(ll_vs, ll_tmp) V_tmp, P_tmp, ll_tmp = vh.forwards_viterbi_hap_low_mem_rescaling( - n, m, H_vs, s, e_vs, r + n=n, + m=m, + H=H_vs, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, + ) + path_tmp = vh.backwards_viterbi_hap(m=m, V_last=V_tmp, P=P_tmp) + ll_check = vh.path_ll_hap( + n=n, + m=m, + H=H_vs, + path=path_tmp, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, ) - path_tmp = vh.backwards_viterbi_hap(m, V_tmp, P_tmp) - ll_check = vh.path_ll_hap(n, m, H_vs, path_tmp, s, e_vs, r) self.assertAllClose(ll_tmp, ll_check) self.assertAllClose(ll_vs, ll_tmp) V_tmp, P_tmp, ll_tmp = vh.forwards_viterbi_hap_lower_mem_rescaling( - n, m, H_vs, s, e_vs, r + n=n, + m=m, + H=H_vs, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, ) - path_tmp = vh.backwards_viterbi_hap(m, V_tmp, P_tmp) + path_tmp = vh.backwards_viterbi_hap(m=m, V_last=V_tmp, P=P_tmp) ll_check = vh.path_ll_hap(n, m, H_vs, path_tmp, s, e_vs, r) self.assertAllClose(ll_tmp, ll_check) self.assertAllClose(ll_vs, ll_tmp) @@ -69,14 +162,29 @@ def verify(self, ts, scale_mutation_rate, include_ancestors): recombs, ll_tmp, ) = vh.forwards_viterbi_hap_lower_mem_rescaling_no_pointer( - n, m, H_vs, s, e_vs, r + n=n, + m=m, + H=H_vs, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, ) path_tmp = vh.backwards_viterbi_hap_no_pointer( - m, - V_argmaxes_tmp, - nb.typed.List(recombs), + m=m, + V_argmaxes=V_argmaxes_tmp, + recombs=nb.typed.List(recombs), + ) + ll_check = vh.path_ll_hap( + n=n, + m=m, + H=H_vs, + path=path_tmp, + s=s, + e=e_vs, + r=r, + emission_func=emission_func, ) - ll_check = vh.path_ll_hap(n, m, H_vs, path_tmp, s, e_vs, r) self.assertAllClose(ll_tmp, ll_check) self.assertAllClose(ll_vs, ll_tmp)