Add recomb rates

tests/test_nontree_vit_haploid_fixed_switches.py

@@ -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
@@ -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):