From 4059d52d7b0f70beadb58b9b33948d4a376a0490 Mon Sep 17 00:00:00 2001
From: Shing Zhan <shing.zhan@gmail.com>
Date: Wed, 30 Aug 2023 16:54:48 +0100
Subject: [PATCH] WIP

---
 python/tests/test_beagle.py | 61 ++++++++++++++++++++++++++-----------
 1 file changed, 43 insertions(+), 18 deletions(-)

diff --git a/python/tests/test_beagle.py b/python/tests/test_beagle.py
index 3f6adc6a04..1f36c1d09f 100644
--- a/python/tests/test_beagle.py
+++ b/python/tests/test_beagle.py
@@ -315,24 +315,6 @@ def compute_state_probability_matrix(fm, bm, ref_h, query_h, rho, mu):
     return (sm, fwd_hap_probs, bwd_hap_probs)
 
 
-def compute_state_probability_matrix_eqn1(fm, bm, ref_h, query_h, rho, mu):
-    """
-    Implement the HMM forward-backward algorithm following Equation 1 of BB2016.
-
-    This is simpler than faithfully reimplementing the original BEAGLE 4.1 algorithm.
-
-    :param numpy.ndarray fm: Forward probability matrix.
-    :param numpy.ndarray bm: Backward probability matrix.
-    :param numpy.ndarray ref_h: Reference haplotypes.
-    :param numpy.ndarray query_h: One query haplotype.
-    :param numpy.ndarray rho: Switch probabilities.
-    :param numpy.ndarray mu: Mismatch probabilities.
-    :return: HMM state probability matrix.
-    :rtype: numpy.ndarray
-    """
-    pass
-
-
 def interpolate_haplotype_probability_matrix(
     fwd_hap_probs, bwd_hap_probs, genotyped_pos, imputed_pos
 ):
@@ -380,6 +362,49 @@ def interpolate_haplotype_probability_matrix(
     return i_hap_probs
 
 
+def compute_state_probability_matrix_equation1(fm, bm, ref_h, query_h, rho, mu):
+    """
+    Implement the HMM forward-backward algorithm following Equation 1 of BB2016.
+    This is simpler than faithfully reimplementing the original BEAGLE 4.1 algorithm.
+
+    :param numpy.ndarray fm: Forward probability matrix.
+    :param numpy.ndarray bm: Backward probability matrix.
+    :param numpy.ndarray ref_h: Reference haplotypes.
+    :param numpy.ndarray query_h: One query haplotype.
+    :param numpy.ndarray rho: Switch probabilities.
+    :param numpy.ndarray mu: Mismatch probabilities.
+    :return: HMM state probability matrix.
+    :rtype: numpy.ndarray
+    """
+    pass
+
+
+def interpolate_allele_probabilities_equation1(sm, ref_h, genotyped_pos, imputed_pos):
+    """
+    Compute the interpolated allele probabilities following Equation 1 of BB2016.
+
+    Assume all biallelic sites.
+
+    :param numpy.ndarray sm: HMM state probability matrix.
+    :param numpy.ndarray ref_h: Reference haplotypes.
+    :param numpy.ndarray genotyped_pos: Site positions at genotyped markers.
+    :param numpy.ndarray imputed_pos: Site positions at imputed markers.
+    :return: Interpolated allele probabilities.
+    :rtype: numpy.ndarray
+    """
+    h = ref_h.shape[1]
+    # m = len(genotyped_pos)
+    x = len(imputed_pos)
+    p = np.array((x, 2), dtype=np.float64)
+    weights = get_weights(genotyped_pos, imputed_pos)
+    # Compute probabilities of allele a at imputed markers
+    a = 0
+    for i in np.arange(x):
+        for j in np.arange(h):
+            p[i, a] = weights[i] * sm[i, j] + (1 - weights[i]) * sm[i + 1, j]
+    return p
+
+
 def get_map_alleles(i_hap_probs):
     """
     Compute the maximum a posteriori alleles at the imputed markers of a query haplotype.