Add wrapper function

szhan · Aug 26, 2023 · c311b19 · c311b19
1 parent 6fa5f3d
commit c311b19
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/python/tests/test_beagle.py b/python/tests/test_beagle.py
@@ -344,3 +344,20 @@ def get_map_alleles(i_hap_probs):
         allele_probs[i, 0] = np.argmax(i_hap_probs[i, :])
         allele_probs[i, 1] = np.argmax(i_hap_probs[i, :])
     return allele_probs
+
+
+def run_beagle(ref_h, query_h, genotyped_pos, imputed_pos):
+    """
+    Run the BEAGLE 4.1 imputation algorithm.
+    """
+    mu = get_mismatch_prob()
+    rho = get_switch_prob()
+    fm = compute_forward_probability_matrix(ref_h, query_h, rho, mu)
+    bm = compute_backward_probability_matrix(ref_h, query_h, rho, mu)
+    _, fwd_hap_probs, bwd_hap_probs = compute_state_probability_matrix(
+        fm, bm, ref_h, query_h, rho, mu
+    )
+    i_hap_probs = compute_interpolated_haplotype_matrix(
+        fwd_hap_probs, bwd_hap_probs, genotyped_pos, imputed_pos
+    )
+    get_map_alleles(i_hap_probs)