Skip to content

Commit

Permalink
Add recomb rates
Browse files Browse the repository at this point in the history
  • Loading branch information
szhan committed Jul 6, 2024
1 parent 22a2e86 commit 741d203
Showing 1 changed file with 38 additions and 10 deletions.
48 changes: 38 additions & 10 deletions tests/test_nontree_vit_haploid_fixed_switches.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,16 +10,40 @@

class TestNonTreeViterbiHaploidFixedSwitches(lsbase.ViterbiAlgorithmBase):
def get_examples_pars():
# One fixed switch.
# Start
# Middle
# End
# Two fixed switches.
#
# Set the recombination rate to zero at all positions
# except the sites of the fixed switches.
# Check the paths that switching only occur there.
pass
num_sites = 10

# Set constant mutation probability.
mu = np.zeros(num_sites, dtype=np.float64) + 1e-4

# Set recombination probabilities for different sets of fixed switches.
recomb_rates_no_switch = np.zeros(num_sites, dtype=np.float64)
recomb_rates_one_switch_start = np.zeros_like(recomb_rates_no_switch)
recomb_rates_one_switch_start[0] = 1e-2
recomb_rates_one_switch_middle = np.zeros_like(recomb_rates_no_switch)
recomb_rates_one_switch_middle[num_sites / 2] = 1e-2
recomb_rates_one_switch_end = np.zeros_like(recomb_rates_no_switch)
recomb_rates_one_switch_end[-1] = 1e-2
recomb_rates_two_switches = np.zeros_like(recomb_rates_no_switch)
recomb_rates_two_switches[3] = 1e-2
recomb_rates_two_switches[6] = 1e-2
recomb_rates_arr = [
recomb_rates_no_switch,
recomb_rates_one_switch_start,
recomb_rates_one_switch_middle,
recomb_rates_one_switch_end,
recomb_rates_two_switches,
]

# Expected paths
path_no_switch = None
path_one_switch_start = None
path_one_switch_middle = None
path_one_switch_end = None
path_two_switches = None

for r in recomb_rates_arr:
yield mu, r


def verify(self, ts, scale_mutation_rate, include_ancestors):
ploidy = 1
Expand Down Expand Up @@ -215,25 +239,29 @@ def verify(self, ts, scale_mutation_rate, include_ancestors):
self.assertAllClose(ll_tmp, ll_check)
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, include_ancestors)


@pytest.mark.parametrize("num_samples", [8, 16, 32])
@pytest.mark.parametrize("scale_mutation_rate", [True, False])
@pytest.mark.parametrize("include_ancestors", [True, False])
def test_ts_simple(self, num_samples, scale_mutation_rate, include_ancestors):
ts = self.get_ts_simple(num_samples)
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_simple_n8_high_recomb(self, scale_mutation_rate, include_ancestors):
ts = self.get_ts_simple_n8_high_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_larger(self, scale_mutation_rate, include_ancestors):
Expand Down

0 comments on commit 741d203

Please sign in to comment.