small issue with Tajimas D tests

Author: petrelharp

python/tests/test_tree_stats.py

@@ -668,10 +668,12 @@ def test_short_sequence_length(self):
 
     @pytest.mark.slow
     def test_wright_fisher_slow(self, wf_fixture_slow):
-        self.verify(wf_fixture_slow)
+        _, ts = wf_fixture_slow
+        self.verify(ts)
 
     def test_wright_fisher(self, wf_fixture):
-        self.verify(wf_fixture)
+        _, ts = wf_fixture
+        self.verify(ts)
 
     def test_empty_ts(self):
         tables = tskit.TableCollection(1.0)
@@ -827,13 +829,13 @@ def wf_fixture_sims():
 )
 def wf_fixture(wf_fixture_sims, request):
     """
-    A collection of small Wright-Fisher simulations.
+    A collection of small Wright-Fisher simulations. The name is returned for
+    debugging purposes.
     """
-    ts = msprime.sim_mutations(
-        wf_fixture_sims[request.param], rate=0.05, random_seed=1234
-    )
+    name = request.param
+    ts = msprime.sim_mutations(wf_fixture_sims[name], rate=0.05, random_seed=1234)
     assert ts.num_mutations > 0
-    return ts
+    return name, ts
 
 
 @pytest.fixture(params=["unsimplified", "initial_generation"], scope="session")
@@ -842,11 +844,10 @@ def wf_fixture_slow(wf_fixture_sims, request):
     A few more small Wright-Fisher simulations. Despite the name, in total
     they take about the same time for tests together as wf_fixture.
     """
-    ts = msprime.sim_mutations(
-        wf_fixture_sims[request.param], rate=0.05, random_seed=1234
-    )
+    name = request.param
+    ts = msprime.sim_mutations(wf_fixture_sims[name], rate=0.05, random_seed=1234)
     assert ts.num_mutations > 0
-    return ts
+    return name, ts
 
 
 @pytest.fixture(scope="session")
@@ -1107,12 +1108,12 @@ def test_many_trees_sequence_length_infinite_sites(
             self.verify(ts)
 
     def test_wright_fisher(self, wf_fixture):
-        ts = wf_fixture
+        _, ts = wf_fixture
         assert ts.num_sites > 0
         self.verify(ts)
 
     def test_wright_fisher_slow(self, wf_fixture_slow):
-        ts = wf_fixture_slow
+        _, ts = wf_fixture_slow
         assert ts.num_sites > 0
         self.verify(ts)
 
@@ -1130,7 +1131,7 @@ def example_sample_sets(ts, min_size=1):
     number of sample sets returned in each example must be at least min_size
     """
     samples = ts.samples()
-    np.random.shuffle(samples)
+    np.random.shuffle(samples)  # <-- no seed
     splits = np.array_split(samples, min_size)
     yield splits
     yield [[s] for s in samples]
@@ -1748,6 +1749,12 @@ def verify_persite_tajimas_d(self, ts, sample_sets):
             sigma1 = ts.Tajimas_D(sample_sets, windows=windows, mode=self.mode)
             sigma2 = site_tajimas_d(ts, sample_sets, windows=windows)
             assert sigma1.shape == sigma2.shape
+            # floating point error can lead in strange cases to
+            # +/-inf in our implementation here and nan in the ts version
+            f1 = np.isfinite(sigma1)
+            f2 = np.isfinite(sigma2)
+            assert np.all(f1 == f2)
+            sigma2[~f2] = np.nan
             self.assertArrayAlmostEqual(sigma1, sigma2)