Refactor site sums

python/tests/beagle_numba.py

@@ -585,10 +585,9 @@ def interpolate_allele_probs(
                 if max(sum_probs_a_k, sum_probs_a_kM1) > _MIN_THRESHOLD:
                     probs[i, a] += w * sum_probs_a_kM1
                     probs[i, a] += (1 - w) * sum_probs_a_k
-    # Rescale probabilities.
-    # TODO: Check if this is necessary. Could this be a subtle source of error?
-    assert np.all(np.sum(probs, axis=1) > 0), "Sum of allele probabilities is not > 0."
-    probs_rescaled = probs / np.sum(probs, axis=1)[:, np.newaxis]
+    site_sums = np.sum(probs, axis=1)
+    assert np.all(site_sums > 0), "Some site sums of allele probabilities is <= 0."
+    probs_rescaled = probs / site_sums[:, np.newaxis]
     if return_weights:
         return (probs_rescaled, weights)
     return (probs_rescaled, None)