Merge pull request #220 from tskit-dev/fst

Fst implementation (and doc fixes)

Author: petrelharp

docs/stats.rst

@@ -140,7 +140,7 @@ in units of crossovers and mutations per base pair, respectively).
 .. _sec_general_stats_sample_sets:
 
 ***********************
-Sample sets and indices
+Sample sets and indexes
 ***********************
 
 Many standard population genetics statistics
@@ -164,40 +164,35 @@ and want to compute **all** fourty-five pairwise divergences?
 You could call ``divergence`` fourty-five times, but this would be tedious
 and also inefficient, because the allele frequencies for one population
 gets used in computing many of those values.
-So, statistics that take a ``sample_sets`` argument also take an ``indices`` argument,
+So, statistics that take a ``sample_sets`` argument also take an ``indexes`` argument,
 which for a statistic that operates on ``k`` sample sets will be a list of ``k``-tuples.
-If ``indices`` is a length ``n`` list of ``k``-tuples,
+If ``indexes`` is a length ``n`` list of ``k``-tuples,
 then the output will have ``n`` columns,
-and if ``indices[j]`` is a tuple ``(i0, ..., ik)``,
+and if ``indexes[j]`` is a tuple ``(i0, ..., ik)``,
 then the ``j``-th column will contain values of the statistic computed on
 ``(sample_sets[i0], sample_sets[i1], ..., sample_sets[ik])``.
 
-To recap: ``indices`` must be a list of tuples, each of length ``k``,
+To recap: ``indexes`` must be a list of tuples, each of length ``k``,
 of integers between ``0`` and ``len(sample_sets) - 1``.
 The appropriate value of ``k`` depends on the statistic.
 
 Here are some additional special cases:
 
-``indices = None``
+``indexes = None``
    If the statistic takes ``k`` inputs for ``k > 1``,
    and there are exactly ``k`` lists in ``sample_sets``,
    then this will compute just one statistic, and is equivalent to passing
-   ``indices = [(0, 1, ..., k-1)]``.
+   ``indexes = [(0, 1, ..., k-1)]``.
    If there are not exactly ``k`` sample sets, this will throw an error.
 
-``k=1`` does not allow ``indices``:
+``k=1`` does not allow ``indexes``:
    Statistics that operate on one sample set at a time (i.e., ``k=1``)
    do **not** take the ``indexes`` argument,
    and instead just return the value of the statistic separately for each of ``sample_sets``
    in the order they are given.
-   (This would be equivalent to passing ``indices = [[0], [1], ..., [len(sample_sets)]]``,
+   (This would be equivalent to passing ``indexes = [[0], [1], ..., [len(sample_sets)]]``,
    were that allowed.)
 
-``stat_type = "node"`` does not allow ``indices``:
-   Since node statistics output one value per node (unlike the other types, which output
-   something summed across all nodes), it is an error to specify ``indices`` when computing
-   a node statistic (consequently, you need to have exactly ``k`` sample sets).
-
 
 .. _sec_general_stats_output:
 
@@ -216,7 +211,6 @@ from ``windows[i]`` to ``windows[i + 1]`` (including the left but not the right
    The output is a two-dimensional array,
    with columns corresponding to the different statistics computed: ``out[i, j]`` is the ``j``-th statistic
    in the ``i``-th window.
-   If the statistic takes an ``indices`` argument, then ``out[i, j]`` has the statistic computed with ``indices[j]``.
 
 ``mode="node"``
    The output is a three-dimensional array,

python/tests/test_tree_stats.py

@@ -125,7 +125,7 @@ def path_length(tr, x, y):
 
 @contextlib.contextmanager
 def suppress_division_by_zero_warning():
-    with np.errstate(invalid='ignore'):
+    with np.errstate(invalid='ignore', divide='ignore'):
         yield
 
 
@@ -830,8 +830,6 @@ def wrapped_summary_func(x):
             self.assertEqual(sigma1.shape, sigma4.shape)
             self.assertArrayAlmostEqual(sigma1, sigma2)
             self.assertArrayAlmostEqual(sigma1, sigma3)
-            # print("computed", sigma1)
-            # print("definition", sigma4)
             self.assertArrayAlmostEqual(sigma1, sigma4)
 
 
@@ -1301,6 +1299,68 @@ class TestSiteDivergence(TestDivergence, MutatedTopologyExamplesMixin):
     mode = "site"
 
 
+############################################
+# Fst
+############################################
+
+def single_site_Fst(ts, sample_sets, indexes):
+    """
+    Compute single-site Fst, which between two groups with frequencies p and q is
+      1 - 2 * (p (1-p) + q(1-q)) / ( p(1-p) + q(1-q) + p(1-q) + q(1-p) )
+    or in the multiallelic case, replacing p(1-p) with the sum over alleles of p(1-p),
+    and adjusted for sampling without replacement.
+    """
+    # TODO: what to do in this case?
+    if ts.num_sites == 0:
+        out = np.array([np.repeat(np.nan, len(indexes))])
+        return out
+    out = np.zeros((ts.num_sites, len(indexes)))
+    samples = ts.samples()
+    for j, v in enumerate(ts.variants()):
+        for i, (ix, iy) in enumerate(indexes):
+            g = v.genotypes
+            X = sample_sets[ix]
+            Y = sample_sets[iy]
+            gX = [a for k, a in zip(samples, g) if k in X]
+            gY = [a for k, a in zip(samples, g) if k in Y]
+            nX = len(X)
+            nY = len(Y)
+            dX = dY = dXY = 0
+            for a in set(g):
+                fX = np.sum(gX == a)
+                fY = np.sum(gY == a)
+                with suppress_division_by_zero_warning():
+                    dX += fX * (nX - fX) / (nX * (nX - 1))
+                    dY += fY * (nY - fY) / (nY * (nY - 1))
+                    dXY += (fX * (nY - fY) + (nX - fX) * fY) / (2 * nX * nY)
+            with suppress_division_by_zero_warning():
+                out[j][i] = 1 - 2 * (dX + dY) / (dX + dY + 2 * dXY)
+    return out
+
+
+class TestFst(StatsTestCase, TwoWaySampleSetStatsMixin):
+
+    # Derived classes define this to get a specific stats mode.
+    mode = None
+
+    def verify(self, ts):
+        # only check per-site
+        for sample_sets in example_sample_sets(ts, min_size=2):
+            for indexes in example_sample_set_index_pairs(sample_sets):
+                self.verify_persite_Fst(ts, sample_sets, indexes)
+
+    def verify_persite_Fst(self, ts, sample_sets, indexes):
+        sigma1 = ts.Fst(sample_sets, indexes=indexes, windows="sites",
+                        mode=self.mode, span_normalise=False)
+        sigma2 = single_site_Fst(ts, sample_sets, indexes)
+        self.assertEqual(sigma1.shape, sigma2.shape)
+        self.assertArrayAlmostEqual(sigma1, sigma2)
+
+
+class TestSiteFst(TestFst, MutatedTopologyExamplesMixin):
+    mode = "site"
+
+
 ############################################
 # Y2
 ############################################
@@ -2289,6 +2349,11 @@ class TestGeneralStatInterface(StatsTestCase):
     Tests for the basic interface for general_stats.
     """
 
+    def get_tree_sequence(self):
+        ts = msprime.simulate(10, recombination_rate=2,
+                              mutation_rate=2, random_seed=1)
+        return ts
+
     def test_default_mode(self):
         ts = msprime.simulate(10, recombination_rate=1, random_seed=2)
         W = np.ones((ts.num_samples, 2))
@@ -2303,6 +2368,15 @@ def test_bad_mode(self):
             with self.assertRaises(ValueError):
                 ts.general_stat(W, lambda x: x, mode=bad_mode)
 
+    def test_bad_window_strings(self):
+        ts = self.get_tree_sequence()
+        with self.assertRaises(ValueError):
+            ts.diversity([list(ts.samples())], mode="site", windows="abc")
+        with self.assertRaises(ValueError):
+            ts.diversity([list(ts.samples())], mode="site", windows="")
+        with self.assertRaises(ValueError):
+            ts.diversity([list(ts.samples())], mode="tree", windows="abc")
+
 
 class TestGeneralBranchStats(StatsTestCase):
     """