Implement adaptive thresholding

python/tests/beagle_numba.py

@@ -482,30 +482,33 @@ def interpolate_allele_probs(
     x = ref_h.shape[0]
     h = ref_h.shape[1]
     # Set threshold to set trivial probabilities to zero.
-    _THRESHOLD = 0
-    if use_threshold:
-        _THRESHOLD = min(0.005, 1 / h)
+    _MIN_THRESHOLD = 0
     weights, right_idx = get_weights(typed_pos, untyped_pos, typed_cm, untyped_cm)
     probs = np.zeros((x, len(alleles)), dtype=np.float64)
     for i in range(x):
         k = right_idx[i]
         w = weights[i]
-        for j in range(h):
-            for a in alleles:
+        for a in alleles:
+            if use_threshold:
+                # TODO: Check whether this is implemented correctly
+                # Threshold based on "the number of subsets in the partition Am of H".
+                threshold_Am = np.sum(ref_h[i, :] == a)
+                _MIN_THRESHOLD = min(0.005, 1 / threshold_Am)
+            for j in range(h):
                 if ref_h[i, j] == a:
                     if k == 0:
                         assert w == 1.0, "Weight should be 1.0."
-                        if sm[k, j] >= _THRESHOLD:
+                        if sm[k, j] >= _MIN_THRESHOLD:
                             probs[i, a] += sm[k, j]
                     elif k == m:
                         # FIXME: It's unclear how BEAGLE handles this case.
                         assert w == 0.0, "Weight should be 0.0."
-                        if sm[k - 1, j] >= _THRESHOLD:
+                        if sm[k - 1, j] >= _MIN_THRESHOLD:
                             probs[i, a] += sm[k - 1, j]
                     else:
-                        if sm[k - 1, j] >= _THRESHOLD:
+                        if sm[k - 1, j] >= _MIN_THRESHOLD:
                             probs[i, a] += sm[k - 1, j]
-                        if sm[k - 1, j] >= _THRESHOLD or sm[k, j] >= _THRESHOLD:
+                        if sm[k - 1, j] >= _MIN_THRESHOLD or sm[k, j] >= _MIN_THRESHOLD:
                             probs[i, a] += w * sm[k - 1, j]
                             probs[i, a] += (1 - w) * sm[k, j]
     # Rescale probabilities.