Fix bugs

python/tests/beagle.py

@@ -429,10 +429,10 @@ def interpolate_allele_probabilities(sm, ref_h, genotyped_pos, imputed_pos):
     Compute the interpolated allele probabilities at imputed markers
     of a query haplotype following Equation 1 of BB2016.
 
-    Assuming all biallelic sites, the interpolated allele probability
-    matrix is of size (x, 2), where x is the number of imputed markers.
+    Assuming all biallelic sites, the interpolated allele probability matrix
+    is of size (x, 2), where x is the number of imputed markers.
 
-    This function takes the output of `_compute_state_probability_matrix`.
+    This function takes the output of `compute_state_probability_matrix`.
 
     Note that this function takes:
     1. HMM state probability matrix across genotyped markers (size of m).
@@ -449,7 +449,7 @@ def interpolate_allele_probabilities(sm, ref_h, genotyped_pos, imputed_pos):
     m = len(genotyped_pos)
     x = len(imputed_pos)
     assert sm.shape == (m, h)
-    assert ref_h.shape == (m + x, h)
+    assert ref_h.shape == (x, h)
     genotyped_cm = convert_to_genetic_map_position(genotyped_pos)
     imputed_cm = convert_to_genetic_map_position(imputed_pos)
     weights = get_weights(genotyped_pos, imputed_pos)
@@ -555,9 +555,7 @@ def run_beagle(ref_h, query_h, pos):
     )
     assert bm.shape == (m, h)
     # Compute HMM state probability matrix at genotyped markers
-    sm = compute_state_probability_matrix(
-        fm, bm, ref_h_genotyped, query_h_genotyped, rho, mu
-    )
+    sm = compute_state_probability_matrix(fm, bm, ref_h_genotyped, query_h_genotyped)
     assert sm.shape == (m, h)
     # Subset the reference haplotypes to imputed markers
     ref_h_imputed = ref_h[imputed_pos_idx, :]