Avoid divide by zero error

Author: tomwhite

sgkit/stats/popgen.py

@@ -46,8 +46,10 @@ def diversity(
     n_pairs = an * (an - 1) / 2
     n_same = (ac * (ac - 1) / 2).sum(axis=2)
     n_diff = n_pairs - n_same
-    pi = n_diff / n_pairs
-    pi_sum = pi.sum(axis=0)
+    # replace zeros to avoid divide by zero error
+    n_pairs_na = n_pairs.where(n_pairs != 0)
+    pi = n_diff / n_pairs_na
+    pi_sum = pi.sum(axis=0, skipna=False)
     return pi_sum  # type: ignore[no-any-return]
 
 
@@ -98,7 +100,7 @@ def Fst(ds: Dataset, allele_counts: Hashable = "cohort_allele_count",) -> DataAr
     DataArray
          fst value between the two cohorts.
     """
-    total_div = diversity(ds, allele_counts).sum()
+    total_div = diversity(ds, allele_counts).sum(skipna=False)
     gs = divergence(ds, allele_counts)
     den = total_div + 2 * gs  # type: ignore[operator]
     fst = 1 - (2 * total_div / den)

sgkit/tests/test_popgen.py

@@ -60,7 +60,7 @@ def test_divergence(size):
     np.testing.assert_allclose(div, ts_div)
 
 
-@pytest.mark.parametrize("size", [10, 100])
+@pytest.mark.parametrize("size", [2, 3, 10, 100])
 def test_Fst(size):
     ts = msprime.simulate(size, length=100, mutation_rate=0.05, random_seed=42)
     subset_1 = ts.samples()[: ts.num_samples // 2]