Handle cases when denominator is too small

python/tests/beagle_numba.py

@@ -749,15 +749,17 @@ def compute_estimated_allelic_r_squared(gt_probs):
     :return: Estimated allelic R-squared.
     :rtype: float
     """
+    _MIN_R2_DEN = 1e-8
     n = len(gt_probs)  # Number of individuals
     z = np.argmax(gt_probs, axis=1)  # Most likely imputed genotypes
     u = gt_probs[:, 1] + 2 * gt_probs[:, 2]  # E[X | y_i]
     w = gt_probs[:, 1] + 4 * gt_probs[:, 2]  # E[X^2 | y_i]
     t_1 = (np.sum(z * u) - (1 / n) * (np.sum(u) * np.sum(z))) ** 2
-    # TODO: Properly handle cases where terms 2 and/or 3 are 0.
     t_2 = np.sum(w) - (1 / n) * np.sum(u) ** 2
     t_3 = np.sum(z**2) - (1 / n) * np.sum(z) ** 2
-    est_allelic_rsq = -1 if t_2 * t_3 == 0 else t_1 / (t_2 * t_3)
+    den = t_2 * t_3
+    # Minimum of allelic R^2 is zero. See BEAGLE code.
+    est_allelic_rsq = 0 if den < _MIN_R2_DEN else t_1 / den
     return est_allelic_rsq