Repository: tskit-dev/tsdate

Issue:
"Invalid parent parameters" error using variational gamma
When running the following code, I get an error. I have attached the trees file so we can debug

```python
import tskit
import tsdate
ts = tskit.load("debug.trees")
tsdate.date(ts, mutation_rate=1e-8, Ne=1e4, progress=True, method="variational_gamma")
```

```
File ~/Library/jupyterlab-desktop/jlab_server/lib/python3.10/site-packages/tsdate/core.py:1050, in ExpectationPropagation.iterate(self, iter_num, progress)
   1048 if iter_num:
   1049     desc = f"EP (iter {iter_num + 1:>2}, leafwards)"
-> 1050 self.propagate(
   1051     edges=self.edges_by_child_desc(grouped=False), desc=desc, progress=progress
   1052 )

File ~/Library/jupyterlab-desktop/jlab_server/lib/python3.10/site-packages/tsdate/core.py:1022, in ExpectationPropagation.propagate(self, edges, desc, progress)
   1012 child_cavity = self.lik.ratio(
   1013     self.posterior[edge.child], self.child_message[edge.id]
   1014 )
   1015 # Get the target posterior: that is, the cavity multiplied by the
   1016 # edge likelihood, and projected onto a gamma distribution via
   1017 # moment matching.
   1018 (
   1019     norm_const,
   1020     self.posterior[edge.parent],
   1021     self.posterior[edge.child],
-> 1022 ) = approx.gamma_projection(*parent_cavity, *child_cavity, *edge_lik)
   1023 # Get the messages: that is, the multiplicative difference between
   1024 # the target and cavity posteriors. This only involves updating the
   1025 # variational parameters for the parent and child on the edge.
   1026 self.parent_message[edge.id] = self.lik.ratio(
   1027     self.posterior[edge.parent], parent_cavity
   1028 )

File ~/Library/jupyterlab-desktop/jlab_server/lib/python3.10/site-packages/tsdate/approx.py:114, in sufficient_statistics()
     95 @numba.njit(
     96     "UniTuple(float64, 5)(float64, float64, float64, float64, float64, float64)"
     97 )
     98 def sufficient_statistics(a_i, b_i, a_j, b_j, y_ij, mu_ij):
     99     """
    100     Calculate gamma sufficient statistics for the PDF proportional to
    101     :math:`Ga(t_j | a_j, b_j) Ga(t_i | a_i, b_i) Po(y_{ij} |
   (...)
    112     :return: normalizing constant, E[t_i], E[log t_i], E[t_j], E[log t_j]
    113     """
--> 114     assert a_i > 0 and b_i >= 0, "Invalid parent parameters"
    115     assert a_j > 0 and b_j >= 0, "Invalid child parameters"
    116     assert y_ij >= 0 and mu_ij > 0, "Invalid edge parameters"

AssertionError: Invalid parent parameters
```

[debug.trees.zip](https://github.com/tskit-dev/tsdate/files/12013054/debug.trees.zip)


Discussion:

Patch:
diff --git a/tsdate/approx.py b/tsdate/approx.py
index c5fd295..0db0a87 100644
--- a/tsdate/approx.py
+++ b/tsdate/approx.py
@@ -111,9 +111,6 @@ def sufficient_statistics(a_i, b_i, a_j, b_j, y_ij, mu_ij):
 
     :return: normalizing constant, E[t_i], E[log t_i], E[t_j], E[log t_j]
     """
-    assert a_i > 0 and b_i >= 0, "Invalid parent parameters"
-    assert a_j > 0 and b_j >= 0, "Invalid child parameters"
-    assert y_ij >= 0 and mu_ij > 0, "Invalid edge parameters"
 
     a = a_i + a_j + y_ij
     b = a_j
@@ -124,9 +121,6 @@ def sufficient_statistics(a_i, b_i, a_j, b_j, y_ij, mu_ij):
         a_i, b_i, a_j, b_j, y_ij, mu_ij
     )
 
-    if sign_f <= 0:
-        raise hypergeo.Invalid2F1("Singular hypergeometric function")
-
     logconst = (
         log_f + hypergeo._betaln(y_ij + 1, b) + hypergeo._gammaln(a) - a * np.log(t)
     )
@@ -142,6 +136,10 @@ def sufficient_statistics(a_i, b_i, a_j, b_j, y_ij, mu_ij):
         - hypergeo._digamma(c)
     )
 
+    # check that Jensen's inequality holds
+    assert np.log(t_i) > ln_t_i
+    assert np.log(t_j) > ln_t_j
+
     return logconst, t_i, ln_t_i, t_j, ln_t_j
 
 
diff --git a/tsdate/hypergeo.py b/tsdate/hypergeo.py
index 5bd396c..e3e14ea 100644
--- a/tsdate/hypergeo.py
+++ b/tsdate/hypergeo.py
@@ -29,8 +29,9 @@ import numba
 import numpy as np
 from numba.extending import get_cython_function_address
 
+# TODO: these are reasonable defaults, but could
+# be made settable via a control dict
 _HYP2F1_TOL = np.sqrt(np.finfo(np.float64).eps)
-_HYP2F1_CHECK = np.sqrt(_HYP2F1_TOL)
 _HYP2F1_MAXTERM = int(1e6)
 
 _PTR = ctypes.POINTER
@@ -115,7 +116,7 @@ def _is_valid_2f1(f1, f2, a, b, c, z):
     See Eq. 6 in https://doi.org/10.1016/j.cpc.2007.11.007
     """
     if z == 0.0:
-        return np.abs(f1 - a * b / c) < _HYP2F1_CHECK
+        return np.abs(f1 - a * b / c) < _HYP2F1_TOL
     u = c - (a + b + 1) * z
     v = a * b
     w = z * (1 - z)
@@ -124,7 +125,7 @@ def _is_valid_2f1(f1, f2, a, b, c, z):
         numer = np.abs(u * f1 - v)
     else:
         numer = np.abs(f2 + u / w * f1 - v / w)
-    return numer / denom < _HYP2F1_CHECK
+    return numer / denom < _HYP2F1_TOL
 
 
 @numba.njit("UniTuple(float64, 7)(float64, float64, float64, float64)")
@@ -255,7 +256,7 @@ def _hyp2f1_recurrence(a, b, c, z):
 
 
 @numba.njit(
-    "UniTuple(float64, 6)(float64, float64, float64, float64, float64, float64)"
+    "UniTuple(float64, 7)(float64, float64, float64, float64, float64, float64)"
 )
 def _hyp2f1_dlmf1583_first(a_i, b_i, a_j, b_j, y, mu):
     """
@@ -287,7 +288,7 @@ def _hyp2f1_dlmf1583_first(a_i, b_i, a_j, b_j, y, mu):
     )
 
     # 2F1(a, -y; c; z) via backwards recurrence
-    val, sign, da, _, dc, dz, _ = _hyp2f1_recurrence(a, y, c, z)
+    val, sign, da, _, dc, dz, d2z = _hyp2f1_recurrence(a, y, c, z)
 
     # map gradient to parameters
     da_i = dc - _digamma(a_i + a_j) + _digamma(a_i)
@@ -295,13 +296,18 @@ def _hyp2f1_dlmf1583_first(a_i, b_i, a_j, b_j, y, mu):
     db_i = dz / (b_j - mu) + a_j / (mu + b_i)
     db_j = dz * (1 - z) / (b_j - mu) - a_j / s / (mu + b_i)
 
+    # needed to verify result
+    d2b_j = (1 - z) / (b_j - mu) ** 2 * (d2z * (1 - z) - 2 * dz * (1 + a_j)) + (
+        1 + a_j
+    ) * a_j / (b_j - mu) ** 2
+
     val += scale
 
-    return val, sign, da_i, db_i, da_j, db_j
+    return val, sign, da_i, db_i, da_j, db_j, d2b_j
 
 
 @numba.njit(
-    "UniTuple(float64, 6)(float64, float64, float64, float64, float64, float64)"
+    "UniTuple(float64, 7)(float64, float64, float64, float64, float64, float64)"
 )
 def _hyp2f1_dlmf1583_second(a_i, b_i, a_j, b_j, y, mu):
     """
@@ -320,7 +326,7 @@ def _hyp2f1_dlmf1583_second(a_i, b_i, a_j, b_j, y, mu):
     )
 
     # 2F1(a, y+1; c; z) via series expansion
-    val, sign, da, _, dc, dz, _ = _hyp2f1_taylor_series(a, y + 1, c, z)
+    val, sign, da, _, dc, dz, d2z = _hyp2f1_taylor_series(a, y + 1, c, z)
 
     # map gradient to parameters
     da_i = da + np.log(z) + dc + _digamma(a_i) - _digamma(a_i + y + 1)
@@ -328,10 +334,16 @@ def _hyp2f1_dlmf1583_second(a_i, b_i, a_j, b_j, y, mu):
     db_i = (1 - z) * (dz + a / z) / (b_i + b_j)
     db_j = -z * (dz + a / z) / (b_i + b_j)
 
+    # needed to verify result
+    d2b_j = (
+        z / (b_i + b_j) ** 2 * (d2z * z + 2 * dz * (1 + a))
+        + a * (1 + a) / (b_i + b_j) ** 2
+    )
+
     sign *= (-1) ** (y + 1)
     val += scale
 
-    return val, sign, da_i, db_i, da_j, db_j
+    return val, sign, da_i, db_i, da_j, db_j, d2b_j
 
 
 @numba.njit(
@@ -345,18 +357,14 @@ def _hyp2f1_dlmf1583(a_i, b_i, a_j, b_j, y, mu):
     assert 0 <= mu <= b_j
     assert y >= 0 and y % 1.0 == 0.0
 
-    f_1, s_1, da_i_1, db_i_1, da_j_1, db_j_1 = _hyp2f1_dlmf1583_first(
+    f_1, s_1, da_i_1, db_i_1, da_j_1, db_j_1, d2b_j_1 = _hyp2f1_dlmf1583_first(
         a_i, b_i, a_j, b_j, y, mu
     )
 
-    f_2, s_2, da_i_2, db_i_2, da_j_2, db_j_2 = _hyp2f1_dlmf1583_second(
+    f_2, s_2, da_i_2, db_i_2, da_j_2, db_j_2, d2b_j_2 = _hyp2f1_dlmf1583_second(
         a_i, b_i, a_j, b_j, y, mu
     )
 
-    if np.abs(f_1 - f_2) < _HYP2F1_TOL:
-        # TODO: detect a priori if this will occur
-        raise Invalid2F1("Singular hypergeometric function")
-
     f_0 = max(f_1, f_2)
     f_1 = np.exp(f_1 - f_0) * s_1
     f_2 = np.exp(f_2 - f_0) * s_2
@@ -366,10 +374,22 @@ def _hyp2f1_dlmf1583(a_i, b_i, a_j, b_j, y, mu):
     db_i = (db_i_1 * f_1 + db_i_2 * f_2) / f
     da_j = (da_j_1 * f_1 + da_j_2 * f_2) / f
     db_j = (db_j_1 * f_1 + db_j_2 * f_2) / f
+    d2b_j = (d2b_j_1 * f_1 + d2b_j_2 * f_2) / f
 
     sign = np.sign(f)
     val = np.log(np.abs(f)) + f_0
 
+    # use first/second derivatives to check that result is non-singular
+    dz = -db_j * (mu + b_i)
+    d2z = d2b_j * (mu + b_i) ** 2
+    if (
+        not _is_valid_2f1(
+            dz, d2z, a_j, a_i + a_j + y, a_j + y + 1, (mu - b_j) / (mu + b_i)
+        )
+        or sign <= 0
+    ):
+        raise Invalid2F1("Hypergeometric series is singular")
+
     return val, sign, da_i, db_i, da_j, db_j
 
 


Test patch:
diff --git a/tests/test_hypergeo.py b/tests/test_hypergeo.py
index 86afdfa..ba248bd 100644
--- a/tests/test_hypergeo.py
+++ b/tests/test_hypergeo.py
@@ -240,3 +240,28 @@ class TestTransforms:
         offset = self._2f1_validate(*pars)
         check = self._2f1_grad_validate(*pars, offset=offset)
         assert np.allclose(grad, check)
+
+
+@pytest.mark.parametrize(
+    "pars",
+    [
+        # taken from examples in issues tsdate/286, tsdate/289
+        [1.104, 0.0001125, 118.1396, 0.009052, 1.0, 0.001404],
+        [2.7481, 0.001221, 344.94083, 0.02329, 3.0, 0.00026624],
+    ],
+)
+class TestSingular2F1:
+    """
+    Test detection of cases where 2F1 is close to singular and DLMF 15.8.3
+    suffers from catastrophic cancellation: in these cases, use DLMF 15.8.1
+    even though it takes much longer to converge.
+    """
+
+    def test_dlmf1583_throws_exception(self, pars):
+        with pytest.raises(Exception, match="is singular"):
+            hypergeo._hyp2f1_dlmf1583(*pars)
+
+    def test_exception_uses_dlmf1581(self, pars):
+        v1, *_ = hypergeo._hyp2f1(*pars)
+        v2, *_ = hypergeo._hyp2f1_dlmf1581(*pars)
+        assert np.isclose(v1, v2)