Minor edit

szhan · Mar 11, 2024 · 663b5db · 663b5db
1 parent 8d5126b
commit 663b5db
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py
@@ -537,7 +537,7 @@ def interpolate_allele_probs(
     See 'setFirstAlleleProbs', 'setAlleleProbs', and 'setLastAlleleProbs'
     in 'LSHapBaum.java' in BEAGLE 4.1 source code.
 
-    :param numpy.ndarray sm: Hidden state probability matrix at genotyped positions.
+    :param numpy.ndarray state_mat: State probability matrix at genotyped positions.
     :param numpy.ndarray ref_h: Reference haplotypes subsetted to ungenotyped positions.
     :param numpy.ndarray pos_typed: Physical positions of genotyped markers (bp).
     :param numpy.ndarray pos_untyped: Physical positions of ungenotyped markers (bp).
@@ -566,20 +566,22 @@ def interpolate_allele_probs(
                 # Threshold based on "the number of subsets in the partition Am of H".
                 threshold_Am = np.sum(is_a_in_ref_h)
                 _MIN_THRESHOLD = min(0.005, 1 / threshold_Am)
-            sum_probs_a_k = np.sum(state_mat[k, is_a_in_ref_h])
-            sum_probs_a_kM1 = np.sum(state_mat[k - 1, is_a_in_ref_h])
             if k == 0:
                 # See 'setFirstAlleleProbs' in 'LSHapBaum.java'.
                 assert w == 1.0, "Weight should be 1.0."
+                sum_probs_a_k = np.sum(state_mat[k, is_a_in_ref_h])
                 if sum_probs_a_k > _MIN_THRESHOLD:
                     probs[i, a] += sum_probs_a_k
             elif k == m:
                 # See 'setLastAlleleProbs' in 'LSHapBaum.java'.
                 assert w == 0.0, "Weight should be 0.0."
+                sum_probs_a_kM1 = np.sum(state_mat[k - 1, is_a_in_ref_h])
                 if sum_probs_a_kM1 > _MIN_THRESHOLD:
                     probs[i, a] += sum_probs_a_kM1
             else:
                 # See 'setAlleleProbs' in 'LSHapBaum.java'.
+                sum_probs_a_k = np.sum(state_mat[k, is_a_in_ref_h])
+                sum_probs_a_kM1 = np.sum(state_mat[k - 1, is_a_in_ref_h])
                 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