Add dimension stripping to derived stats.

Author: jeromekelleher

python/tests/test_tree_stats.py

@@ -1719,12 +1719,9 @@ def verify_persite_Fst(self, ts, sample_sets, indexes):
         self.assertArrayAlmostEqual(sigma1, sigma2)
 
 
-class FstInterfaceMixin(StatsTestCase):
+class FstInterfaceMixin(object):
 
-    # Since Fst is defined using diversity and divergence, we don't seriously
-    # test it for correctness, and only test the interface.
-
-    def verify_interface(self):
+    def test_interface(self):
         ts = msprime.simulate(10, mutation_rate=0.0)
         sample_sets = [[0, 1, 2], [6, 7], [4]]
         with self.assertRaises(ValueError):
@@ -1735,32 +1732,23 @@ def verify_interface(self):
             ts.Fst(sample_sets, indexes=[(0, 1), (0, 20)])
         sigma1 = ts.Fst(sample_sets, indexes=[(0, 1)], mode=self.mode)
         sigma2 = ts.Fst(sample_sets, indexes=[(0, 1), (0, 2), (1, 2)], mode=self.mode)
-        if self.mode == "node":
-            self.assertArrayAlmostEqual(sigma1[:, :, 0], sigma2[:, :, 0])
-        else:
-            self.assertArrayAlmostEqual(sigma1[:, 0], sigma2[:, 0])
+        self.assertArrayAlmostEqual(sigma1[..., 0], sigma2[..., 0])
 
 
 class TestSiteFst(TestFst, MutatedTopologyExamplesMixin, FstInterfaceMixin):
     mode = "site"
 
-    def test_interface(self):
-        self.verify_interface()
 
+# Since Fst is defined using diversity and divergence, we don't seriously
+# test it for correctness for node and branch, and only test the interface.
 
-class TestNodeFst(FstInterfaceMixin):
+class TestNodeFst(StatsTestCase, FstInterfaceMixin):
     mode = "node"
 
-    def test_interface(self):
-        self.verify_interface()
 
-
-class TestBranchFst(FstInterfaceMixin):
+class TestBranchFst(StatsTestCase, FstInterfaceMixin):
     mode = "node"
 
-    def test_interface(self):
-        self.verify_interface()
-
 
 ############################################
 # Y2
@@ -5105,56 +5093,67 @@ def test_one_way_stat_default_windows(self):
         # We're adding on the *last* dimension, so must reshape
         self.assertArrayEqual(x.reshape(ts.num_nodes, 1), y)
 
-    def test_one_way_stat_windows(self):
-        ts = self.get_example_ts()
+    def verify_one_way_stat_windows(self, ts, method):
         L = ts.sequence_length
         N = ts.num_nodes
 
-        windows = [0, L / 4, L / 2, L]
+        windows = [0, L / 4, L / 2, 0.75 * L, L]
         A = ts.samples()[:6]
         B = ts.samples()[6:]
         for mode in ["site", "branch"]:
-            x = ts.diversity([A, B], windows=windows, mode=mode)
-            # Three windows, 2 sets.
-            self.assertEqual(x.shape, (3, 2))
+            x = method([A, B], windows=windows, mode=mode)
+            # Four windows, 2 sets.
+            self.assertEqual(x.shape, (4, 2))
 
-            x = ts.diversity([A], windows=windows, mode=mode)
-            # Three windows, 1 sets.
-            self.assertEqual(x.shape, (3, 1))
+            x = method([A], windows=windows, mode=mode)
+            # Four windows, 1 sets.
+            self.assertEqual(x.shape, (4, 1))
 
-            x = ts.diversity(A, windows=windows, mode=mode)
+            x = method(A, windows=windows, mode=mode)
             # Dropping the outer list removes the last dimension
-            self.assertEqual(x.shape, (3, ))
+            self.assertEqual(x.shape, (4, ))
 
-            x = ts.diversity(windows=windows, mode=mode)
+            x = method(windows=windows, mode=mode)
             # Default returns this for all samples
-            self.assertEqual(x.shape, (3, ))
-            y = ts.diversity(ts.samples(), windows=windows, mode=mode)
+            self.assertEqual(x.shape, (4, ))
+            y = method(ts.samples(), windows=windows, mode=mode)
             self.assertArrayEqual(x, y)
 
         mode = "node"
-        x = ts.diversity([A, B], windows=windows, mode=mode)
-        # Three windows, N nodes and 2 sets.
-        self.assertEqual(x.shape, (3, N, 2))
+        x = method([A, B], windows=windows, mode=mode)
+        # Four windows, N nodes and 2 sets.
+        self.assertEqual(x.shape, (4, N, 2))
 
-        x = ts.diversity([A], windows=windows, mode=mode)
-        # Three windows, N nodes and 1 set.
-        self.assertEqual(x.shape, (3, N, 1))
+        x = method([A], windows=windows, mode=mode)
+        # Four windows, N nodes and 1 set.
+        self.assertEqual(x.shape, (4, N, 1))
 
-        x = ts.diversity(A, windows=windows, mode=mode)
+        x = method(A, windows=windows, mode=mode)
         # Drop the outer list, so we lose the last dimension
-        self.assertEqual(x.shape, (3, N))
+        self.assertEqual(x.shape, (4, N))
 
-        x = ts.diversity(windows=windows, mode=mode)
+        x = method(windows=windows, mode=mode)
         # The default sample sets also drops the last dimension
-        self.assertEqual(x.shape, (3, N))
+        self.assertEqual(x.shape, (4, N))
 
         self.assertEqual(ts.num_trees, 1)
         # In this example, we know that the trees are all the same so check this
         # for sanity.
         self.assertArrayEqual(x[0], x[1])
         self.assertArrayEqual(x[0], x[2])
 
+    def test_diversity_windows(self):
+        ts = self.get_example_ts()
+        self.verify_one_way_stat_windows(ts, ts.diversity)
+
+    def test_Tajimas_D_windows(self):
+        ts = self.get_example_ts()
+        self.verify_one_way_stat_windows(ts, ts.Tajimas_D)
+
+    def test_segregating_sites_windows(self):
+        ts = self.get_example_ts()
+        self.verify_one_way_stat_windows(ts, ts.segregating_sites)
+
     def test_two_way_stat_default_windows(self):
         ts = self.get_example_ts()
         # Use divergence as the example one-way stat.
@@ -5179,47 +5178,45 @@ def test_two_way_stat_default_windows(self):
         # We're adding on the *last* dimension, so must reshape
         self.assertArrayEqual(x.reshape(ts.num_nodes, 1), y)
 
-    def test_two_way_stat_windows(self):
-
-        ts = self.get_example_ts()
+    def verify_two_way_stat_windows(self, ts, method):
         L = ts.sequence_length
         N = ts.num_nodes
 
         windows = [0, L / 4, L / 2, L]
         A = ts.samples()[:7]
         B = ts.samples()[7:]
         for mode in ["site", "branch"]:
-            x = ts.divergence(
+            x = method(
                 [A, B, A], indexes=[[0, 1], [0, 2]], windows=windows, mode=mode)
             # Three windows, 2 pairs
             self.assertEqual(x.shape, (3, 2))
 
-            x = ts.divergence([A, B], indexes=[[0, 1]], windows=windows, mode=mode)
+            x = method([A, B], indexes=[[0, 1]], windows=windows, mode=mode)
             # Three windows, 1 pair
             self.assertEqual(x.shape, (3, 1))
 
-            x = ts.divergence([A, B], indexes=[0, 1], windows=windows, mode=mode)
+            x = method([A, B], indexes=[0, 1], windows=windows, mode=mode)
             # Dropping the outer list removes the last dimension
             self.assertEqual(x.shape, (3, ))
 
-            y = ts.divergence([A, B], windows=windows, mode=mode)
+            y = method([A, B], windows=windows, mode=mode)
             self.assertEqual(y.shape, (3, ))
             self.assertArrayEqual(x, y)
 
         mode = "node"
-        x = ts.divergence([A, B], indexes=[[0, 1], [0, 1]], windows=windows, mode=mode)
+        x = method([A, B], indexes=[[0, 1], [0, 1]], windows=windows, mode=mode)
         # Three windows, N nodes and 2 pairs
         self.assertEqual(x.shape, (3, N, 2))
 
-        x = ts.divergence([A, B], indexes=[[0, 1]], windows=windows, mode=mode)
+        x = method([A, B], indexes=[[0, 1]], windows=windows, mode=mode)
         # Three windows, N nodes and 1 pairs
         self.assertEqual(x.shape, (3, N, 1))
 
-        x = ts.divergence([A, B], indexes=[0, 1], windows=windows, mode=mode)
+        x = method([A, B], indexes=[0, 1], windows=windows, mode=mode)
         # Drop the outer list, so we lose the last dimension
         self.assertEqual(x.shape, (3, N))
 
-        x = ts.divergence([A, B], windows=windows, mode=mode)
+        x = method([A, B], windows=windows, mode=mode)
         # The default sample sets also drops the last dimension
         self.assertEqual(x.shape, (3, N))
 
@@ -5229,8 +5226,19 @@ def test_two_way_stat_windows(self):
         self.assertArrayEqual(x[0], x[1])
         self.assertArrayEqual(x[0], x[2])
 
-    def test_three_way_stat_windows(self):
+    def test_divergence_windows(self):
         ts = self.get_example_ts()
+        self.verify_two_way_stat_windows(ts, ts.divergence)
+
+    def test_Fst_windows(self):
+        ts = self.get_example_ts()
+        self.verify_two_way_stat_windows(ts, ts.Fst)
+
+    def test_f2_windows(self):
+        ts = self.get_example_ts()
+        self.verify_two_way_stat_windows(ts, ts.f2)
+
+    def verify_three_way_stat_windows(self, ts, method):
         L = ts.sequence_length
         N = ts.num_nodes
 
@@ -5239,37 +5247,37 @@ def test_three_way_stat_windows(self):
         B = ts.samples()[2: 4]
         C = ts.samples()[4:]
         for mode in ["site", "branch"]:
-            x = ts.Y3(
+            x = method(
                 [A, B, C], indexes=[[0, 1, 2], [0, 2, 1]], windows=windows, mode=mode)
             # Three windows, 2 triple
             self.assertEqual(x.shape, (3, 2))
 
-            x = ts.Y3([A, B, C], indexes=[[0, 1, 2]], windows=windows, mode=mode)
+            x = method([A, B, C], indexes=[[0, 1, 2]], windows=windows, mode=mode)
             # Three windows, 1 triple
             self.assertEqual(x.shape, (3, 1))
 
-            x = ts.Y3([A, B, C], indexes=[0, 1, 2], windows=windows, mode=mode)
+            x = method([A, B, C], indexes=[0, 1, 2], windows=windows, mode=mode)
             # Dropping the outer list removes the last dimension
             self.assertEqual(x.shape, (3, ))
 
-            y = ts.Y3([A, B, C], windows=windows, mode=mode)
+            y = method([A, B, C], windows=windows, mode=mode)
             self.assertEqual(y.shape, (3, ))
             self.assertArrayEqual(x, y)
 
         mode = "node"
-        x = ts.Y3([A, B, C], indexes=[[0, 1, 2], [0, 2, 1]], windows=windows, mode=mode)
+        x = method([A, B, C], indexes=[[0, 1, 2], [0, 2, 1]], windows=windows, mode=mode)
         # Three windows, N nodes and 2 triples
         self.assertEqual(x.shape, (3, N, 2))
 
-        x = ts.Y3([A, B, C], indexes=[[0, 1, 2]], windows=windows, mode=mode)
+        x = method([A, B, C], indexes=[[0, 1, 2]], windows=windows, mode=mode)
         # Three windows, N nodes and 1 triples
         self.assertEqual(x.shape, (3, N, 1))
 
-        x = ts.Y3([A, B, C], indexes=[0, 1, 2], windows=windows, mode=mode)
+        x = method([A, B, C], indexes=[0, 1, 2], windows=windows, mode=mode)
         # Drop the outer list, so we lose the last dimension
         self.assertEqual(x.shape, (3, N))
 
-        x = ts.Y3([A, B, C], windows=windows, mode=mode)
+        x = method([A, B, C], windows=windows, mode=mode)
         # The default sample sets also drops the last dimension
         self.assertEqual(x.shape, (3, N))
 
@@ -5279,6 +5287,13 @@ def test_three_way_stat_windows(self):
         self.assertArrayEqual(x[0], x[1])
         self.assertArrayEqual(x[0], x[2])
 
+    def test_Y3_windows(self):
+        ts = self.get_example_ts()
+        self.verify_three_way_stat_windows(ts, ts.Y3)
+
+    def test_f3_windows(self):
+        ts = self.get_example_ts()
+        self.verify_three_way_stat_windows(ts, ts.f3)
 
 ############################################
 # Old code where stats are defined within type
@@ -5287,6 +5302,7 @@ def test_three_way_stat_windows(self):
 # The only thing left to port is the SFS code.
 ############################################
 
+
 class PythonBranchStatCalculator(object):
     """
     Python implementations of various ("tree") branch-length statistics -