Update API

szhan · Jul 1, 2024 · 692b477 · 692b477
1 parent 0044769
commit 692b477
Showing 1 changed file with 65 additions and 44 deletions.
diff --git a/lshmm/api.py b/lshmm/api.py
@@ -223,24 +223,34 @@ def forwards(
     )
 
     if ploidy == 1:
-        forward_function = forwards_ls_hap
+        (
+            forward_array,
+            normalisation_factor_from_forward,
+            log_lik,
+        ) = forwards_ls_hap(
+            num_ref_haps,
+            num_sites,
+            ref_panel_checked,
+            query_checked,
+            emission_matrix,
+            prob_recombination,
+            norm=normalise,
+            emission_func=core.get_emission_probability_haploid,
+        )
     else:
-        forward_function = forward_ls_dip_loop
-
-    (
-        forward_array,
-        normalisation_factor_from_forward,
-        log_lik,
-    ) = forward_function(
-        num_ref_haps,
-        num_sites,
-        ref_panel_checked,
-        query_checked,
-        emission_matrix,
-        prob_recombination,
-        norm=normalise,
-        emission_func=core.get_emission_probability_haploid,
-    )
+        (
+            forward_array,
+            normalisation_factor_from_forward,
+            log_lik,
+        ) = forward_ls_dip_loop(
+            num_ref_haps,
+            num_sites,
+            ref_panel_checked,
+            query_checked,
+            emission_matrix,
+            prob_recombination,
+            norm=normalise,
+        )
 
     return forward_array, normalisation_factor_from_forward, log_lik
 
@@ -268,20 +278,26 @@ def backwards(
     )
 
     if ploidy == 1:
-        backward_function = backwards_ls_hap
+        backwards_array = backwards_ls_hap(
+            num_ref_haps,
+            num_sites,
+            ref_panel_checked,
+            query_checked,
+            emission_matrix,
+            normalisation_factor_from_forward,
+            prob_recombination,
+            emission_func=core.get_emission_probability_haploid,
+        )
     else:
-        backward_function = backward_ls_dip_loop
-
-    backwards_array = backward_function(
-        num_ref_haps,
-        num_sites,
-        ref_panel_checked,
-        query_checked,
-        emission_matrix,
-        normalisation_factor_from_forward,
-        prob_recombination,
-        emission_func=core.get_emission_probability_haploid,
-    )
+        backwards_array = backward_ls_dip_loop(
+            num_ref_haps,
+            num_sites,
+            ref_panel_checked,
+            query_checked,
+            emission_matrix,
+            normalisation_factor_from_forward,
+            prob_recombination,
+        )
 
     return backwards_array
 
@@ -326,7 +342,6 @@ def viterbi(
             query_checked,
             emission_matrix,
             prob_recombination,
-            emission_func=core.get_emission_probability_diploid,
         )
         unphased_path = backwards_viterbi_dip(num_sites, V, P)
         best_path = get_phased_path(num_ref_haps, unphased_path)
@@ -357,19 +372,25 @@ def path_loglik(
     )
 
     if ploidy == 1:
-        path_ll_function = path_ll_hap
+        log_lik = path_ll_hap(
+            num_ref_haps,
+            num_sites,
+            ref_panel_checked,
+            path,
+            query_checked,
+            emission_matrix,
+            prob_recombination,
+            emission_func=core.get_emission_probability_haploid,
+        )
     else:
-        path_ll_function = path_ll_dip
-
-    log_lik = path_ll_function(
-        num_ref_haps,
-        num_sites,
-        ref_panel_checked,
-        path,
-        query_checked,
-        emission_matrix,
-        prob_recombination,
-        emission_func=core.get_emission_probability_haploid,
-    )
+        log_lik = path_ll_dip(
+            num_ref_haps,
+            num_sites,
+            ref_panel_checked,
+            path,
+            query_checked,
+            emission_matrix,
+            prob_recombination,
+        )
 
     return log_lik