From 670e2baae84e1762c22a273ac97a03154c8a96e9 Mon Sep 17 00:00:00 2001
From: Shing Zhan <shing.zhan@gmail.com>
Date: Wed, 13 Sep 2023 13:44:11 +0100
Subject: [PATCH] Assume all biallelic sites

---
 python/tests/beagle.py       | 6 ++----
 python/tests/beagle_numba.py | 6 ++----
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/python/tests/beagle.py b/python/tests/beagle.py
index 3c4b2c5bfc..3a298e9193 100644
--- 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,7 +476,6 @@ 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
     """
@@ -486,6 +483,7 @@ def interpolate_allele_probabilities(
     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
index 66b3621e23..c0937f8840 100644
--- 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