Assume all biallelic sites

szhan · Sep 13, 2023 · 670e2ba · 670e2ba
1 parent 6020892
commit 670e2ba
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 8 deletions.
diff --git a/python/tests/beagle.py b/python/tests/beagle.py
@@ -458,9 +458,7 @@ def compute_state_probability_matrix(fm, bm, ref_h, query_h):
     return sm
 
 
-def interpolate_allele_probabilities(
-    sm, ref_h, genotyped_pos, imputed_pos, alleles=(0, 1)
-):
+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.
@@ -478,14 +476,14 @@ def interpolate_allele_probabilities(
     :param numpy.ndarray ref_h: Reference haplotypes subsetted to imputed markers.
     :param numpy.ndarray genotyped_pos: Site positions at genotyped markers.
     :param numpy.ndarray imputed_pos: Site positions at imputed markers.
-    :param list alleles: Distince alleles in the haplotypes.
     :return: Interpolated allele probabilities.
     :rtype: numpy.ndarray
     """
     assert not np.any(sm < 0), "HMM state probability matrix has negative values."
     assert not np.any(np.isnan(sm)), "HMM state probability matrix has NaN values."
     m = sm.shape[0] - 1
     h = sm.shape[1]
+    alleles = [0, 1]  # Assume all biallelic sites
     assert np.all(
         np.isin(ref_h, alleles)
     ), f"Reference haplotypes have alleles absent in {alleles}."

diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py
@@ -263,9 +263,7 @@ def compute_state_probability_matrix(fm, bm, ref_h, query_h):
 
 
 @njit
-def interpolate_allele_probabilities(
-    sm, ref_h, genotyped_pos, imputed_pos, alleles=(0, 1)
-):
+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.
@@ -283,12 +281,12 @@ def interpolate_allele_probabilities(
     :param numpy.ndarray ref_h: Reference haplotypes subsetted to imputed markers.
     :param numpy.ndarray genotyped_pos: Site positions at genotyped markers.
     :param numpy.ndarray imputed_pos: Site positions at imputed markers.
-    :param list alleles: Distince alleles in the haplotypes.
     :return: Interpolated allele probabilities.
     :rtype: numpy.ndarray
     """
     m = sm.shape[0] - 1
     h = sm.shape[1]
+    alleles = [0, 1]  # Assume all biallelic sites
     assert m == len(genotyped_pos)
     x = len(imputed_pos)
     assert (x, h) == ref_h.shape