Implement reference haplotype segment indexing

Author: szhan

python/tests/test_beagle.py

@@ -278,15 +278,32 @@ def compute_state_probability_matrix(fm, bm, ref_h, query_h, rho, mu):
     assert np.all(np.isin(np.unique(ref_h), [0, 1]))
     assert np.all(np.isin(np.unique(query_h), [-1, 0, 1]))
     sm = np.zeros((m, h), dtype=np.float64)  # HMM state probability matrix
-    fwd_hap_probs = np.zeros((m, 2), dtype=np.float64)
-    bwd_hap_probs = np.zeros((m, 2), dtype=np.float64)
+
+    def _get_ref_hap_seg(a, b):
+        """
+        Assuming all biallelic sites, get the index of a reference haplotype segment (i)
+        defined by the alleles a and b at two adjacent genotyped markers.
+
+        #i, a, b
+        0, 0, 0
+        1, 1, 0
+        2, 0, 1
+        3, 1, 1
+
+        See https://github.com/tskit-dev/tskit/issues/2802#issuecomment-1698767169
+        """
+        ref_hap_idx = a * 1 + b * 2
+        return ref_hap_idx
+
+    fwd_hap_probs = np.zeros((m, 4), dtype=np.float64)
+    bwd_hap_probs = np.zeros((m, 4), dtype=np.float64)
     for i in np.arange(m - 1, 0, -1):
         for j in np.arange(h):
             sm[i, j] = fm[i, j] * bm[i, j]
-            ref_allele_mP1 = ref_h[i + 1, j]
-            ref_allele_m = ref_h[i, j]
-            fwd_hap_probs[i, ref_allele_mP1] += sm[i, j]
-            bwd_hap_probs[i, ref_allele_m] += sm[i, j]
+            ref_hap_seg_mP1 = _get_ref_hap_seg(ref_h[i, j], ref_h[i + 1, j])
+            ref_hap_seg_m = _get_ref_hap_seg(ref_h[i - 1, j], ref_h[i, j])
+            fwd_hap_probs[i, ref_hap_seg_mP1] += sm[i, j]
+            bwd_hap_probs[i, ref_hap_seg_m] += sm[i, j]
         sum_fwd_hap_probs = np.sum(fwd_hap_probs[i, :])
         fwd_hap_probs[i, :] /= sum_fwd_hap_probs
         bwd_hap_probs[i, :] /= sum_fwd_hap_probs  # Sum of forward hap probs