From 6fb8bab7762d529a6acf46f2716f3f47da7ab63d Mon Sep 17 00:00:00 2001 From: szhan Date: Tue, 18 Jun 2024 13:55:07 +0100 Subject: [PATCH] Split tests by algorithm --- ..._api_diploid.py => test_api_fb_diploid.py} | 99 ------------ ..._api_haploid.py => test_api_fb_haploid.py} | 94 ------------ ...llelic.py => test_api_fb_haploid_multi.py} | 55 ------- tests/test_api_vit_diploid.py | 104 +++++++++++++ tests/test_api_vit_haploid.py | 99 ++++++++++++ tests/test_api_vit_haploid_multi.py | 60 ++++++++ ..._diploid.py => test_nontree_fb_diploid.py} | 137 +---------------- tests/test_nontree_fb_haploid.py | 93 ++++++++++++ tests/test_nontree_vit_diploid.py | 143 ++++++++++++++++++ ...haploid.py => test_nontree_vit_haploid.py} | 87 +---------- 10 files changed, 501 insertions(+), 470 deletions(-) rename tests/{test_api_diploid.py => test_api_fb_diploid.py} (53%) rename tests/{test_api_haploid.py => test_api_fb_haploid.py} (53%) rename tests/{test_API_multiallelic.py => test_api_fb_haploid_multi.py} (55%) create mode 100644 tests/test_api_vit_diploid.py create mode 100644 tests/test_api_vit_haploid.py create mode 100644 tests/test_api_vit_haploid_multi.py rename tests/{test_nontree_diploid.py => test_nontree_fb_diploid.py} (52%) create mode 100644 tests/test_nontree_fb_haploid.py create mode 100644 tests/test_nontree_vit_diploid.py rename tests/{test_nontree_haploid.py => test_nontree_vit_haploid.py} (61%) diff --git a/tests/test_api_diploid.py b/tests/test_api_fb_diploid.py similarity index 53% rename from tests/test_api_diploid.py rename to tests/test_api_fb_diploid.py index 0588c0f..4e6dfad 100644 --- a/tests/test_api_diploid.py +++ b/tests/test_api_fb_diploid.py @@ -4,7 +4,6 @@ import lshmm as ls import lshmm.core as core import lshmm.fb_diploid as fbd -import lshmm.vit_diploid as vd class TestForwardBackwardDiploid(lsbase.ForwardBackwardAlgorithmBase): @@ -122,101 +121,3 @@ def test_ts_larger(self, scale_mutation_rate, include_ancestors): scale_mutation_rate=scale_mutation_rate, include_ancestors=include_ancestors, ) - - -class TestViterbiDiploid(lsbase.ViterbiAlgorithmBase): - def verify(self, ts, scale_mutation_rate, include_ancestors): - for n, m, H_vs, query, e_vs, r, mu in self.get_examples_pars( - ts, - ploidy=2, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - include_extreme_rates=True, - ): - G_vs = core.convert_haplotypes_to_phased_genotypes(H_vs) - s = core.convert_haplotypes_to_unphased_genotypes(query) - num_alleles = core.get_num_alleles(H_vs, query) - - V_vs, P_vs, ll_vs = vd.forwards_viterbi_dip_low_mem( - n=n, - m=m, - G=G_vs, - s=s, - e=e_vs, - r=r, - ) - path_vs = vd.backwards_viterbi_dip(m=m, V_last=V_vs, P=P_vs) - phased_path_vs = vd.get_phased_path(n=n, path=path_vs) - path, ll = ls.viterbi( - reference_panel=G_vs, - query=s, - num_alleles=num_alleles, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=scale_mutation_rate, - ) - - self.assertAllClose(ll_vs, ll) - self.assertAllClose(phased_path_vs, path) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n10_no_recomb(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n10_no_recomb() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n6(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n6() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n8(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n8() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n8_high_recomb(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n8_high_recomb() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n16(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n16() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_larger(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_custom_pars( - ref_panel_size=45, length=1e5, mean_r=1e-5, mean_mu=1e-5 - ) - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) diff --git a/tests/test_api_haploid.py b/tests/test_api_fb_haploid.py similarity index 53% rename from tests/test_api_haploid.py rename to tests/test_api_fb_haploid.py index 70ca656..536e896 100644 --- a/tests/test_api_haploid.py +++ b/tests/test_api_fb_haploid.py @@ -4,7 +4,6 @@ import lshmm as ls import lshmm.core as core import lshmm.fb_haploid as fbh -import lshmm.vit_haploid as vh class TestForwardBackwardHaploid(lsbase.ForwardBackwardAlgorithmBase): @@ -117,96 +116,3 @@ def test_ts_larger(self, scale_mutation_rate, include_ancestors): scale_mutation_rate=scale_mutation_rate, include_ancestors=include_ancestors, ) - - -class TestViterbiHaploid(lsbase.ViterbiAlgorithmBase): - def verify(self, ts, scale_mutation_rate, include_ancestors): - for n, m, H_vs, s, e_vs, r, mu in self.get_examples_pars( - ts, - ploidy=1, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - include_extreme_rates=True, - ): - num_alleles = core.get_num_alleles(H_vs, s) - V_vs, P_vs, ll_vs = vh.forwards_viterbi_hap_lower_mem_rescaling( - n=n, - m=m, - H=H_vs, - s=s, - e=e_vs, - r=r, - ) - path_vs = vh.backwards_viterbi_hap(m=m, V_last=V_vs, P=P_vs) - path, ll = ls.viterbi( - reference_panel=H_vs, - query=s, - num_alleles=num_alleles, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=scale_mutation_rate, - ) - self.assertAllClose(ll_vs, ll) - self.assertAllClose(path_vs, path) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n10_no_recomb(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n10_no_recomb() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n6(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n6() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n8(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n8() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n8_high_recomb(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n8_high_recomb() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n16(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n16() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_larger(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_custom_pars( - ref_panel_size=46, length=1e5, mean_r=1e-5, mean_mu=1e-5 - ) - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) diff --git a/tests/test_API_multiallelic.py b/tests/test_api_fb_haploid_multi.py similarity index 55% rename from tests/test_API_multiallelic.py rename to tests/test_api_fb_haploid_multi.py index 40aa198..6eb416c 100644 --- a/tests/test_API_multiallelic.py +++ b/tests/test_api_fb_haploid_multi.py @@ -4,7 +4,6 @@ import lshmm as ls import lshmm.core as core import lshmm.fb_haploid as fbh -import lshmm.vit_haploid as vh class TestForwardBackwardHaploid(lsbase.ForwardBackwardAlgorithmBase): @@ -81,57 +80,3 @@ def test_ts_multiallelic_n8(self, scale_mutation_rate, include_ancestors): def test_ts_multiallelic_n16(self, scale_mutation_rate, include_ancestors): ts = self.get_ts_multiallelic_n16() self.verify(ts, scale_mutation_rate, include_ancestors) - - -class TestViterbiHaploid(lsbase.ViterbiAlgorithmBase): - def verify(self, ts, scale_mutation_rate, include_ancestors): - for n, m, H_vs, s, e_vs, r, mu in self.get_examples_pars( - ts, - ploidy=1, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - include_extreme_rates=True, - ): - num_alleles = core.get_num_alleles(H_vs, s) - V_vs, P_vs, ll_vs = vh.forwards_viterbi_hap_lower_mem_rescaling( - n, m, H_vs, s, e_vs, r - ) - path_vs = vh.backwards_viterbi_hap(m, V_vs, P_vs) - path_ll_hap = vh.path_ll_hap(n, m, H_vs, path_vs, s, e_vs, r) - path, ll = ls.viterbi( - reference_panel=H_vs, - query=s, - num_alleles=num_alleles, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=scale_mutation_rate, - ) - self.assertAllClose(ll_vs, ll) - self.assertAllClose(ll_vs, path_ll_hap) - self.assertAllClose(path_vs, path) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_multiallelic_n10_no_recomb( - self, scale_mutation_rate, include_ancestors - ): - ts = self.get_ts_multiallelic_n10_no_recomb() - self.verify(ts, scale_mutation_rate, include_ancestors) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_multiallelic_n6(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_multiallelic_n6() - self.verify(ts, scale_mutation_rate, include_ancestors) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_multiallelic_n8(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_multiallelic_n8() - self.verify(ts, scale_mutation_rate, include_ancestors) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_multiallelic_n16(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_multiallelic_n16() - self.verify(ts, scale_mutation_rate, include_ancestors) diff --git a/tests/test_api_vit_diploid.py b/tests/test_api_vit_diploid.py new file mode 100644 index 0000000..f66715a --- /dev/null +++ b/tests/test_api_vit_diploid.py @@ -0,0 +1,104 @@ +import pytest + +from . import lsbase +import lshmm as ls +import lshmm.core as core +import lshmm.vit_diploid as vd + + +class TestViterbiDiploid(lsbase.ViterbiAlgorithmBase): + def verify(self, ts, scale_mutation_rate, include_ancestors): + for n, m, H_vs, query, e_vs, r, mu in self.get_examples_pars( + ts, + ploidy=2, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + include_extreme_rates=True, + ): + G_vs = core.convert_haplotypes_to_phased_genotypes(H_vs) + s = core.convert_haplotypes_to_unphased_genotypes(query) + num_alleles = core.get_num_alleles(H_vs, query) + + V_vs, P_vs, ll_vs = vd.forwards_viterbi_dip_low_mem( + n=n, + m=m, + G=G_vs, + s=s, + e=e_vs, + r=r, + ) + path_vs = vd.backwards_viterbi_dip(m=m, V_last=V_vs, P=P_vs) + phased_path_vs = vd.get_phased_path(n=n, path=path_vs) + path, ll = ls.viterbi( + reference_panel=G_vs, + query=s, + num_alleles=num_alleles, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=scale_mutation_rate, + ) + + self.assertAllClose(ll_vs, ll) + self.assertAllClose(phased_path_vs, path) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n10_no_recomb(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n10_no_recomb() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n6(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n6() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n8(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n8() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n8_high_recomb(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n8_high_recomb() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n16(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n16() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_larger(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_custom_pars( + ref_panel_size=45, length=1e5, mean_r=1e-5, mean_mu=1e-5 + ) + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) diff --git a/tests/test_api_vit_haploid.py b/tests/test_api_vit_haploid.py new file mode 100644 index 0000000..dd78bfc --- /dev/null +++ b/tests/test_api_vit_haploid.py @@ -0,0 +1,99 @@ +import pytest + +from . import lsbase +import lshmm as ls +import lshmm.core as core +import lshmm.vit_haploid as vh + + +class TestViterbiHaploid(lsbase.ViterbiAlgorithmBase): + def verify(self, ts, scale_mutation_rate, include_ancestors): + for n, m, H_vs, s, e_vs, r, mu in self.get_examples_pars( + ts, + ploidy=1, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + include_extreme_rates=True, + ): + num_alleles = core.get_num_alleles(H_vs, s) + V_vs, P_vs, ll_vs = vh.forwards_viterbi_hap_lower_mem_rescaling( + n=n, + m=m, + H=H_vs, + s=s, + e=e_vs, + r=r, + ) + path_vs = vh.backwards_viterbi_hap(m=m, V_last=V_vs, P=P_vs) + path, ll = ls.viterbi( + reference_panel=H_vs, + query=s, + num_alleles=num_alleles, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=scale_mutation_rate, + ) + self.assertAllClose(ll_vs, ll) + self.assertAllClose(path_vs, path) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n10_no_recomb(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n10_no_recomb() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n6(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n6() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n8(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n8() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n8_high_recomb(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n8_high_recomb() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n16(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n16() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_larger(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_custom_pars( + ref_panel_size=46, length=1e5, mean_r=1e-5, mean_mu=1e-5 + ) + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) diff --git a/tests/test_api_vit_haploid_multi.py b/tests/test_api_vit_haploid_multi.py new file mode 100644 index 0000000..15bbfde --- /dev/null +++ b/tests/test_api_vit_haploid_multi.py @@ -0,0 +1,60 @@ +import pytest + +from . import lsbase +import lshmm as ls +import lshmm.core as core +import lshmm.vit_haploid as vh + + +class TestViterbiHaploid(lsbase.ViterbiAlgorithmBase): + def verify(self, ts, scale_mutation_rate, include_ancestors): + for n, m, H_vs, s, e_vs, r, mu in self.get_examples_pars( + ts, + ploidy=1, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + include_extreme_rates=True, + ): + num_alleles = core.get_num_alleles(H_vs, s) + V_vs, P_vs, ll_vs = vh.forwards_viterbi_hap_lower_mem_rescaling( + n, m, H_vs, s, e_vs, r + ) + path_vs = vh.backwards_viterbi_hap(m, V_vs, P_vs) + path_ll_hap = vh.path_ll_hap(n, m, H_vs, path_vs, s, e_vs, r) + path, ll = ls.viterbi( + reference_panel=H_vs, + query=s, + num_alleles=num_alleles, + prob_recombination=r, + prob_mutation=mu, + scale_mutation_rate=scale_mutation_rate, + ) + self.assertAllClose(ll_vs, ll) + self.assertAllClose(ll_vs, path_ll_hap) + self.assertAllClose(path_vs, path) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_multiallelic_n10_no_recomb( + self, scale_mutation_rate, include_ancestors + ): + ts = self.get_ts_multiallelic_n10_no_recomb() + self.verify(ts, scale_mutation_rate, include_ancestors) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_multiallelic_n6(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_multiallelic_n6() + self.verify(ts, scale_mutation_rate, include_ancestors) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_multiallelic_n8(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_multiallelic_n8() + self.verify(ts, scale_mutation_rate, include_ancestors) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_multiallelic_n16(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_multiallelic_n16() + self.verify(ts, scale_mutation_rate, include_ancestors) diff --git a/tests/test_nontree_diploid.py b/tests/test_nontree_fb_diploid.py similarity index 52% rename from tests/test_nontree_diploid.py rename to tests/test_nontree_fb_diploid.py index 64d077c..f1febfc 100644 --- a/tests/test_nontree_diploid.py +++ b/tests/test_nontree_fb_diploid.py @@ -1,11 +1,11 @@ import pytest + import numpy as np import numba as nb from . import lsbase import lshmm.core as core import lshmm.fb_diploid as fbd -import lshmm.vit_diploid as vd class TestNonTreeForwardBackwardDiploid(lsbase.ForwardBackwardAlgorithmBase): @@ -163,138 +163,3 @@ def test_ts_larger(self, scale_mutation_rate, include_ancestors, normalise): normalise=normalise, include_extreme_rates=include_extreme_rates, ) - - -class TestNonTreeViterbiDiploid(lsbase.ViterbiAlgorithmBase): - def verify(self, ts, scale_mutation_rate, include_ancestors): - for n, m, H_vs, query, e_vs, r, _ in self.get_examples_pars( - ts, - ploidy=2, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - include_extreme_rates=True, - ): - G_vs = core.convert_haplotypes_to_phased_genotypes(H_vs) - s = core.convert_haplotypes_to_unphased_genotypes(query) - - V_vs, P_vs, ll_vs = vd.forwards_viterbi_dip_low_mem(n, m, G_vs, s, e_vs, r) - path_vs = vd.backwards_viterbi_dip(m, V_vs, P_vs) - phased_path_vs = vd.get_phased_path(n, path_vs) - path_ll_vs = vd.path_ll_dip(n, m, G_vs, phased_path_vs, s, e_vs, r) - self.assertAllClose(ll_vs, path_ll_vs) - - ( - V_tmp, - V_argmaxes_tmp, - V_rowcol_maxes_tmp, - V_rowcol_argmaxes_tmp, - recombs_single, - recombs_double, - ll_tmp, - ) = vd.forwards_viterbi_dip_low_mem_no_pointer(n, m, G_vs, s, e_vs, r) - path_tmp = vd.backwards_viterbi_dip_no_pointer( - m, - V_argmaxes_tmp, - V_rowcol_maxes_tmp, - V_rowcol_argmaxes_tmp, - nb.typed.List(recombs_single), - nb.typed.List(recombs_double), - V_tmp, - ) - phased_path_tmp = vd.get_phased_path(n, path_tmp) - path_ll_tmp = vd.path_ll_dip(n, m, G_vs, phased_path_tmp, s, e_vs, r) - self.assertAllClose(ll_tmp, path_ll_tmp) - self.assertAllClose(ll_vs, ll_tmp) - - MAX_NUM_REF_HAPS = 50 - num_ref_haps = H_vs.shape[1] - if num_ref_haps <= MAX_NUM_REF_HAPS: - # Run tests for the naive implementations. - V_tmp, P_tmp, ll_tmp = vd.forwards_viterbi_dip_naive( - n, m, G_vs, s, e_vs, r - ) - path_tmp = vd.backwards_viterbi_dip(m, V_tmp[m - 1, :, :], P_tmp) - phased_path_tmp = vd.get_phased_path(n, path_tmp) - path_ll_tmp = vd.path_ll_dip(n, m, G_vs, phased_path_tmp, s, e_vs, r) - self.assertAllClose(ll_tmp, path_ll_tmp) - self.assertAllClose(ll_vs, ll_tmp) - - V_tmp, P_tmp, ll_tmp = vd.forwards_viterbi_dip_naive_low_mem( - n, m, G_vs, s, e_vs, r - ) - path_tmp = vd.backwards_viterbi_dip(m, V_tmp, P_tmp) - phased_path_tmp = vd.get_phased_path(n, path_tmp) - path_ll_tmp = vd.path_ll_dip(n, m, G_vs, phased_path_tmp, s, e_vs, r) - self.assertAllClose(ll_tmp, path_ll_tmp) - self.assertAllClose(ll_vs, ll_tmp) - - V_tmp, P_tmp, ll_tmp = vd.forwards_viterbi_dip_naive_vec( - n, m, G_vs, s, e_vs, r - ) - path_tmp = vd.backwards_viterbi_dip(m, V_tmp[m - 1, :, :], P_tmp) - phased_path_tmp = vd.get_phased_path(n, path_tmp) - path_ll_tmp = vd.path_ll_dip(n, m, G_vs, phased_path_tmp, s, e_vs, r) - self.assertAllClose(ll_tmp, path_ll_tmp) - self.assertAllClose(ll_vs, ll_tmp) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n10_no_recomb(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n10_no_recomb() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n6(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n6() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n8(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n8() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n8_high_recomb(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n8_high_recomb() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n16(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n16() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_larger(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_custom_pars( - ref_panel_size=30, length=1e5, mean_r=1e-6, mean_mu=1e-5 - ) - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) diff --git a/tests/test_nontree_fb_haploid.py b/tests/test_nontree_fb_haploid.py new file mode 100644 index 0000000..58c7623 --- /dev/null +++ b/tests/test_nontree_fb_haploid.py @@ -0,0 +1,93 @@ +import pytest + +import numpy as np +import numba as nb + +from . import lsbase +import lshmm.core as core +import lshmm.fb_haploid as fbh + + +class TestNonTreeForwardBackwardHaploid(lsbase.ForwardBackwardAlgorithmBase): + def verify(self, ts, scale_mutation_rate, include_ancestors): + for n, m, H_vs, s, e_vs, r, _ in self.get_examples_pars( + ts, + ploidy=1, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + include_extreme_rates=True, + ): + F_vs, c_vs, ll_vs = fbh.forwards_ls_hap(n, m, H_vs, s, e_vs, r, norm=False) + B_vs = fbh.backwards_ls_hap(n, m, H_vs, s, e_vs, c_vs, r) + 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, norm=True + ) + B_tmp = fbh.backwards_ls_hap(n, m, H_vs, s, e_vs, c_tmp, r) + self.assertAllClose(np.sum(F_tmp * B_tmp, 1), np.ones(m)) + self.assertAllClose(ll_vs, ll_tmp) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n10_no_recomb(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n10_no_recomb() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n6(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n6() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n8(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n8() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n8_high_recomb(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n8_high_recomb() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n16(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n16() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_larger(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_custom_pars( + ref_panel_size=45, + length=1e5, + mean_r=1e-5, + mean_mu=1e-5, + ) + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) diff --git a/tests/test_nontree_vit_diploid.py b/tests/test_nontree_vit_diploid.py new file mode 100644 index 0000000..23c39b2 --- /dev/null +++ b/tests/test_nontree_vit_diploid.py @@ -0,0 +1,143 @@ +import pytest + +import numpy as np +import numba as nb + +from . import lsbase +import lshmm.core as core +import lshmm.vit_diploid as vd + + +class TestNonTreeViterbiDiploid(lsbase.ViterbiAlgorithmBase): + def verify(self, ts, scale_mutation_rate, include_ancestors): + for n, m, H_vs, query, e_vs, r, _ in self.get_examples_pars( + ts, + ploidy=2, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + include_extreme_rates=True, + ): + G_vs = core.convert_haplotypes_to_phased_genotypes(H_vs) + s = core.convert_haplotypes_to_unphased_genotypes(query) + + V_vs, P_vs, ll_vs = vd.forwards_viterbi_dip_low_mem(n, m, G_vs, s, e_vs, r) + path_vs = vd.backwards_viterbi_dip(m, V_vs, P_vs) + phased_path_vs = vd.get_phased_path(n, path_vs) + path_ll_vs = vd.path_ll_dip(n, m, G_vs, phased_path_vs, s, e_vs, r) + self.assertAllClose(ll_vs, path_ll_vs) + + ( + V_tmp, + V_argmaxes_tmp, + V_rowcol_maxes_tmp, + V_rowcol_argmaxes_tmp, + recombs_single, + recombs_double, + ll_tmp, + ) = vd.forwards_viterbi_dip_low_mem_no_pointer(n, m, G_vs, s, e_vs, r) + path_tmp = vd.backwards_viterbi_dip_no_pointer( + m, + V_argmaxes_tmp, + V_rowcol_maxes_tmp, + V_rowcol_argmaxes_tmp, + nb.typed.List(recombs_single), + nb.typed.List(recombs_double), + V_tmp, + ) + phased_path_tmp = vd.get_phased_path(n, path_tmp) + path_ll_tmp = vd.path_ll_dip(n, m, G_vs, phased_path_tmp, s, e_vs, r) + self.assertAllClose(ll_tmp, path_ll_tmp) + self.assertAllClose(ll_vs, ll_tmp) + + MAX_NUM_REF_HAPS = 50 + num_ref_haps = H_vs.shape[1] + if num_ref_haps <= MAX_NUM_REF_HAPS: + # Run tests for the naive implementations. + V_tmp, P_tmp, ll_tmp = vd.forwards_viterbi_dip_naive( + n, m, G_vs, s, e_vs, r + ) + path_tmp = vd.backwards_viterbi_dip(m, V_tmp[m - 1, :, :], P_tmp) + phased_path_tmp = vd.get_phased_path(n, path_tmp) + path_ll_tmp = vd.path_ll_dip(n, m, G_vs, phased_path_tmp, s, e_vs, r) + self.assertAllClose(ll_tmp, path_ll_tmp) + self.assertAllClose(ll_vs, ll_tmp) + + V_tmp, P_tmp, ll_tmp = vd.forwards_viterbi_dip_naive_low_mem( + n, m, G_vs, s, e_vs, r + ) + path_tmp = vd.backwards_viterbi_dip(m, V_tmp, P_tmp) + phased_path_tmp = vd.get_phased_path(n, path_tmp) + path_ll_tmp = vd.path_ll_dip(n, m, G_vs, phased_path_tmp, s, e_vs, r) + self.assertAllClose(ll_tmp, path_ll_tmp) + self.assertAllClose(ll_vs, ll_tmp) + + V_tmp, P_tmp, ll_tmp = vd.forwards_viterbi_dip_naive_vec( + n, m, G_vs, s, e_vs, r + ) + path_tmp = vd.backwards_viterbi_dip(m, V_tmp[m - 1, :, :], P_tmp) + phased_path_tmp = vd.get_phased_path(n, path_tmp) + path_ll_tmp = vd.path_ll_dip(n, m, G_vs, phased_path_tmp, s, e_vs, r) + self.assertAllClose(ll_tmp, path_ll_tmp) + self.assertAllClose(ll_vs, ll_tmp) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n10_no_recomb(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n10_no_recomb() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n6(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n6() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n8(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n8() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n8_high_recomb(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n8_high_recomb() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_simple_n16(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_simple_n16() + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) + + @pytest.mark.parametrize("scale_mutation_rate", [True, False]) + @pytest.mark.parametrize("include_ancestors", [True, False]) + def test_ts_larger(self, scale_mutation_rate, include_ancestors): + ts = self.get_ts_custom_pars( + ref_panel_size=30, length=1e5, mean_r=1e-6, mean_mu=1e-5 + ) + self.verify( + ts, + scale_mutation_rate=scale_mutation_rate, + include_ancestors=include_ancestors, + ) diff --git a/tests/test_nontree_haploid.py b/tests/test_nontree_vit_haploid.py similarity index 61% rename from tests/test_nontree_haploid.py rename to tests/test_nontree_vit_haploid.py index 623d0b8..837f59e 100644 --- a/tests/test_nontree_haploid.py +++ b/tests/test_nontree_vit_haploid.py @@ -1,98 +1,13 @@ import pytest + import numpy as np import numba as nb from . import lsbase import lshmm.core as core -import lshmm.fb_haploid as fbh import lshmm.vit_haploid as vh -class TestNonTreeForwardBackwardHaploid(lsbase.ForwardBackwardAlgorithmBase): - def verify(self, ts, scale_mutation_rate, include_ancestors): - for n, m, H_vs, s, e_vs, r, _ in self.get_examples_pars( - ts, - ploidy=1, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - include_extreme_rates=True, - ): - F_vs, c_vs, ll_vs = fbh.forwards_ls_hap(n, m, H_vs, s, e_vs, r, norm=False) - B_vs = fbh.backwards_ls_hap(n, m, H_vs, s, e_vs, c_vs, r) - 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, norm=True - ) - B_tmp = fbh.backwards_ls_hap(n, m, H_vs, s, e_vs, c_tmp, r) - self.assertAllClose(np.sum(F_tmp * B_tmp, 1), np.ones(m)) - self.assertAllClose(ll_vs, ll_tmp) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n10_no_recomb(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n10_no_recomb() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n6(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n6() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n8(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n8() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n8_high_recomb(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n8_high_recomb() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_ts_simple_n16(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_simple_n16() - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - @pytest.mark.parametrize("scale_mutation_rate", [True, False]) - @pytest.mark.parametrize("include_ancestors", [True, False]) - def test_larger(self, scale_mutation_rate, include_ancestors): - ts = self.get_ts_custom_pars( - ref_panel_size=45, - length=1e5, - mean_r=1e-5, - mean_mu=1e-5, - ) - self.verify( - ts, - scale_mutation_rate=scale_mutation_rate, - include_ancestors=include_ancestors, - ) - - class TestNonTreeViterbiHaploid(lsbase.ViterbiAlgorithmBase): def verify(self, ts, scale_mutation_rate, include_ancestors): for n, m, H_vs, s, e_vs, r, _ in self.get_examples_pars(