diff --git a/python/tests/test_beagle.py b/python/tests/test_beagle.py
index 3cdfad1fc5..c0c691d56a 100644
--- a/python/tests/test_beagle.py
+++ b/python/tests/test_beagle.py
@@ -409,22 +409,21 @@ def compute_state_probability_matrix_equation1(fm, bm, ref_h, query_h, rho, mu):
     return sm
 
 
-def interpolate_allele_probabilities_equation1(
-    sm, ref_h, query_h, genotyped_pos, imputed_pos
-):
+def interpolate_allele_probabilities_equation1(sm, ref_h, genotyped_pos, imputed_pos):
     """
-    Compute the interpolated allele probabilities following Equation 1 of BB2016.
+    Compute the interpolated allele probabilities at imputed markers
+    following Equation 1 of BB2016.
 
     This function takes the output of `compute_state_probability_matrix_equation1`.
 
     Assume all biallelic sites.
 
-    Note that this function takes the full reference haplotypes and query haplotype,
-    i.e., not subsetted to genotyped markers.
+    Note that this function takes:
+    1) HMM state probability matrix across genotyped markers (size of m).
+    2) reference haplotypes subsetted to imputed markers (size of x).
 
-    :param numpy.ndarray sm: HMM state probability matrix.
-    :param numpy.ndarray ref_h: Reference haplotypes.
-    :param numpy.ndarray query_h: One query haplotype.
+    :param numpy.ndarray sm: HMM state probability matrix at genotyped markers.
+    :param numpy.ndarray ref_h: Reference haplotypes subsetted to imputed markers.
     :param numpy.ndarray genotyped_pos: Site positions at genotyped markers.
     :param numpy.ndarray imputed_pos: Site positions at imputed markers.
     :return: Interpolated allele probabilities.
@@ -435,7 +434,6 @@ def interpolate_allele_probabilities_equation1(
     x = len(imputed_pos)
     assert sm.shape == (m, h)
     assert ref_h.shape == (m + x, h)
-    assert len(query_h) == m + x
     genotyped_cm = convert_to_genetic_map_position(genotyped_pos)
     imputed_cm = convert_to_genetic_map_position(imputed_pos)
     weights = get_weights(genotyped_pos, imputed_pos)
@@ -445,6 +443,8 @@ def interpolate_allele_probabilities_equation1(
     def _compute_allele_probabilities(a):
         """Equation 1 in BB2016."""
         for i in np.arange(x):
+            # Identify the genotyped markers m and m + 1
+            # that flank the imputed marker i
             if imputed_cm[i] < genotyped_cm[0]:
                 _m = 0
             elif imputed_cm[i] > genotyped_cm[-1]:
@@ -452,8 +452,7 @@ def _compute_allele_probabilities(a):
             else:
                 _m = np.min(np.where(genotyped_cm > imputed_cm[i]))
             for j in np.arange(h):
-                k = 0  # Index wrt reference markers
-                if ref_h[k, j] == a:
+                if ref_h[i, j] == a:
                     p[i, a] += weights[i] * sm[_m, j]
                     p[i, a] += (1 - weights[i]) * sm[_m + 1, j]