Divergence should handle case when all calls for a cohort are missing #402

Author: tomwhite

sgkit/stats/popgen.py

@@ -147,13 +147,14 @@ def _divergence(ac: ArrayLike, out: ArrayLike) -> None:  # pragma: no cover
     for i in range(n_cohorts):
         for j in range(i + 1, n_cohorts):
             n_pairs = an[i] * an[j]
-            n_same = 0
-            for k in range(n_alleles):
-                n_same += ac[i, k] * ac[j, k]
-            n_diff = n_pairs - n_same
-            div = n_diff / n_pairs
-            out[i, j] = div
-            out[j, i] = div
+            if n_pairs != 0.0:
+                n_same = 0
+                for k in range(n_alleles):
+                    n_same += ac[i, k] * ac[j, k]
+                n_diff = n_pairs - n_same
+                div = n_diff / n_pairs
+                out[i, j] = div
+                out[j, i] = div
 
     # calculate the diversity for each cohort
     for i in range(n_cohorts):

sgkit/tests/test_popgen.py

@@ -22,6 +22,8 @@
 )
 from sgkit.window import window
 
+from .test_aggregation import get_dataset
+
 
 def ts_to_dataset(ts, chunks=None, samples=None):
     """
@@ -212,6 +214,17 @@ def test_divergence__windowed_scikit_allel_comparison(sample_size, n_cohorts, ch
     )  # scikit-allel has final window missing
 
 
+def test_divergence__missing_calls():
+    ds = get_dataset(
+        [
+            [[0, 0], [-1, -1], [-1, -1]],  # all of cohort 1 calls are missing
+        ]
+    )
+    ds["sample_cohort"] = xr.DataArray(np.array([0, 1, 1]), dims="samples")
+    ds = divergence(ds)
+    np.testing.assert_equal(ds["stat_divergence"].values[0, 1], np.nan)
+
+
 @pytest.mark.parametrize("sample_size", [2, 3, 10, 100])
 def test_Fst__Hudson(sample_size):
     # scikit-allel can only calculate Fst for pairs of cohorts (populations)