Refactor

Author: szhan

python/tests/test_beagle.py

@@ -228,6 +228,25 @@ def compute_backward_probability_matrix(ref_h, query_h, rho, mu):
     return bm
 
 
+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
+
+    This is a helper function for `compute_state_probability_matrix`.
+    """
+    ref_hap_idx = a * 1 + b * 2
+    return ref_hap_idx
+
+
 def compute_state_probability_matrix(fm, bm, ref_h, query_h, rho, mu):
     """
     Implement the HMM forward-backward algorithm to compute:
@@ -278,23 +297,6 @@ 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
-
-    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):