Clean up

tests/test_API.py

@@ -101,7 +101,7 @@ def _get_num_alleles(ref_haps, query):
 
         # Mixture of random and extremes
         rs = [np.zeros(m) + 0.999, np.zeros(m) + 1e-6, np.random.rand(m)]
-        mus = [np.zeros(m) + 0.33, np.zeros(m) + 1e-6, np.random.rand(m) * 0.33]
+        mus = [np.zeros(m) + 0.2, np.zeros(m) + 1e-6, np.random.rand(m) * 0.2]
 
         for s, r, mu in itertools.product(haplotypes, rs, mus):
             r[0] = 0
@@ -193,29 +193,52 @@ def assertAllClose(self, A, B):
     # Define a bunch of very small tree-sequences for testing a collection of parameters on
     def test_simple_n_10_no_recombination(self):
         ts = msprime.simulate(
-            10, recombination_rate=0, mutation_rate=0.5, random_seed=42
+            samples=10,
+            recombination_rate=0,
+            mutation_rate=0.5,
+            random_seed=42,
         )
         assert ts.num_sites > 3
         self.verify(ts)
 
     def test_simple_n_6(self):
-        ts = msprime.simulate(6, recombination_rate=2, mutation_rate=7, random_seed=42)
+        ts = msprime.simulate(
+            samples=6,
+            recombination_rate=2,
+            mutation_rate=7,
+            random_seed=42,
+        )
         assert ts.num_sites > 5
         self.verify(ts)
 
     def test_simple_n_8(self):
-        ts = msprime.simulate(8, recombination_rate=2, mutation_rate=5, random_seed=42)
+        ts = msprime.simulate(
+            samples=8,
+            recombination_rate=2,
+            mutation_rate=5,
+            random_seed=42,
+        )
         assert ts.num_sites > 5
         self.verify(ts)
 
     def test_simple_n_8_high_recombination(self):
-        ts = msprime.simulate(8, recombination_rate=20, mutation_rate=5, random_seed=42)
+        ts = msprime.simulate(
+            samples=8,
+            recombination_rate=20,
+            mutation_rate=5,
+            random_seed=42,
+        )
         assert ts.num_trees > 15
         assert ts.num_sites > 5
         self.verify(ts)
 
     def test_simple_n_16(self):
-        ts = msprime.simulate(16, recombination_rate=2, mutation_rate=5, random_seed=42)
+        ts = msprime.simulate(
+            samples=16,
+            recombination_rate=2,
+            mutation_rate=5,
+            random_seed=42,
+        )
         assert ts.num_sites > 5
         self.verify(ts)
 

tests/test_API_multiallelic.py

@@ -1,8 +1,9 @@
 import itertools
+import pytest
 
-import msprime
 import numpy as np
-import pytest
+
+import msprime
 import tskit
 
 import lshmm as ls
@@ -16,7 +17,7 @@
 class LSBase:
     """Superclass of Li and Stephens tests."""
 
-    def example_haplotypes(self, ts, num_random=10, seed=42):
+    def example_haplotypes(self, ts, seed=42):
         H = ts.genotype_matrix()
         s = H[:, 0].reshape(1, H.shape[0])
         H = H[:, 1:]
@@ -92,10 +93,6 @@ def _get_num_alleles(ref_haps, query):
         rs = [np.zeros(m) + 0.999, np.zeros(m) + 1e-6, np.random.rand(m)]
         mus = [np.zeros(m) + 0.2, np.zeros(m) + 1e-6, np.random.rand(m) * 0.2]
 
-        e = self.haplotype_emission(
-            mu, m, n_alleles, scale_mutation_based_on_n_alleles=scale_mutation
-        )
-
         for s, r, mu in itertools.product(haplotypes, rs, mus):
             r[0] = 0
             n_alleles = _get_num_alleles(H, s)
@@ -112,11 +109,15 @@ def test_simple_n_10_no_recombination(self):
         ts = msprime.sim_ancestry(
             samples=10,
             recombination_rate=0,
-            random_seed=42,
             sequence_length=10,
             population_size=10000,
+            random_seed=42,
+        )
+        ts = msprime.sim_mutations(
+            ts,
+            rate=1e-5,
+            random_seed=42,
         )
-        ts = msprime.sim_mutations(ts, rate=1e-5, random_seed=42)
         assert ts.num_sites > 3
         self.verify(ts)
 
@@ -128,19 +129,27 @@ def test_simple_n_6(self):
             sequence_length=40,
             population_size=10000,
         )
-        ts = msprime.sim_mutations(ts, rate=1e-3, random_seed=42)
+        ts = msprime.sim_mutations(
+            ts,
+            rate=1e-3,
+            random_seed=42,
+        )
         assert ts.num_sites > 5
         self.verify(ts)
 
     def test_simple_n_8(self):
         ts = msprime.sim_ancestry(
             samples=8,
             recombination_rate=1e-4,
-            random_seed=42,
             sequence_length=20,
             population_size=10000,
+            random_seed=42,
+        )
+        ts = msprime.sim_mutations(
+            ts,
+            rate=1e-4,
+            random_seed=42
         )
-        ts = msprime.sim_mutations(ts, rate=1e-4, random_seed=42)
         assert ts.num_sites > 5
         assert ts.num_trees > 15
         self.verify(ts)
@@ -149,11 +158,15 @@ def test_simple_n_16(self):
         ts = msprime.sim_ancestry(
             samples=16,
             recombination_rate=1e-2,
-            random_seed=42,
             sequence_length=20,
             population_size=10000,
+            random_seed=42,
+        )
+        ts = msprime.sim_mutations(
+            ts,
+            rate=1e-4,
+            random_seed=42,
         )
-        ts = msprime.sim_mutations(ts, rate=1e-4, random_seed=42)
         assert ts.num_sites > 5
         self.verify(ts)
 
@@ -166,7 +179,7 @@ class FBAlgorithmBase(LSBase):
 
 
 class TestMethodsHap(FBAlgorithmBase):
-    """Test that we compute the same likelihoods across all implementations."""
+    """Test that the computed likelihood is the same across all implementations."""
 
     def verify(self, ts):
         for n, m, H_vs, s, e_vs, r, mu in self.example_parameters_haplotypes(ts):