WIP

Author: szhan

python/tests/test_beagle.py

@@ -419,6 +419,9 @@ def interpolate_allele_probabilities_equation1(
 
     Assume all biallelic sites.
 
+    Note that this function takes the full reference haplotypes and query haplotype,
+    i.e., not subsetted to genotyped markers.
+
     :param numpy.ndarray sm: HMM state probability matrix.
     :param numpy.ndarray ref_h: Reference haplotypes.
     :param numpy.ndarray query_h: One query haplotype.
@@ -433,22 +436,26 @@ def interpolate_allele_probabilities_equation1(
     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)
     assert len(weights) == x
     p = np.zeros((x, 2), dtype=np.float64)
 
     def _compute_allele_probabilities(a):
-        """Helper function to compute probability of allele a at imputed markers."""
-        k = 0  # Keep track of imputed marker index
-        # l = 0  # Keep track of genotyped marker index
-        for i in np.arange(m + x):
-            if query_h[i] != -1:
-                continue
+        """Equation 1 in BB2016."""
+        for i in np.arange(x):
+            if imputed_cm[i] < genotyped_cm[0]:
+                _m = 0
+            elif imputed_cm[i] > genotyped_cm[-1]:
+                _m = -2
+            else:
+                _m = np.min(np.where(genotyped_cm > imputed_cm[i]))
             for j in np.arange(h):
-                if ref_h[i, j] == a:
-                    p[k, a] += weights[i] * sm[i, j]
-                    p[k, a] += (1 - weights[i]) * sm[i + 1, j]
-            k += 1
+                k = 0  # Index wrt reference markers
+                if ref_h[k, j] == a:
+                    p[i, a] += weights[i] * sm[_m, j]
+                    p[i, a] += (1 - weights[i]) * sm[_m + 1, j]
 
     _compute_allele_probabilities(a=0)
     _compute_allele_probabilities(a=1)