Pass use_threshold

Author: szhan

python/tests/beagle_numba.py

@@ -623,6 +623,7 @@ def run_interpolation_beagle(
     ne=1e6,
     error_rate=1e-4,
     genetic_map=None,
+    use_threshold=False,
 ):
     """
     Perform a simplified version of the procedure of interpolation-style imputation
@@ -644,6 +645,7 @@ def run_interpolation_beagle(
     :param int ne: Effective population size (default = 1e6).
     :param float error_rate: Allele error rate (default = 1e-4).
     :param GeneticMap genetic_map: Genetic map (default = None).
+    :param bool use_threshold: Set trivial probabilities to 0 if True (default = False).
     :return: Imputed alleles and their probabilities.
     :rtype: tuple(numpy.ndarray, numpy.ndarray)
     """
@@ -694,7 +696,7 @@ def run_interpolation_beagle(
         pos_untyped=pos_untyped,
         cm_typed=cm_typed,
         cm_untyped=cm_untyped,
-        use_threshold=False,
+        use_threshold=use_threshold,
         return_weights=False,
     )
     imputed_alleles, max_allele_probs = get_map_alleles(imputed_allele_probs)
@@ -710,6 +712,7 @@ def run_tsimpute(
     error_rate=1e-4,
     precision=10,
     genetic_map=None,
+    use_threshold=False,
 ):
     """
     Perform interpolation-style imputation, except that the forward and backward
@@ -744,6 +747,7 @@ def run_tsimpute(
     :param float error_rate: Allelic error rate (default = 1e-4).
     :param int precision: Precision for running LS HMM (default = 10).
     :param GeneticMap genetic_map: Genetic map (default = None).
+    :param bool use_threshold: Set trivial probabilities to 0 if True (default = False).
     :return: Imputed alleles and their probabilities.
     :rtype: tuple(numpy.ndarray, numpy.ndarray)
     """
@@ -792,7 +796,7 @@ def run_tsimpute(
         pos_untyped=pos_untyped,
         cm_typed=cm_typed,
         cm_untyped=cm_untyped,
-        use_threshold=False,
+        use_threshold=use_threshold,
         return_weights=False,
     )
     imputed_alleles, max_allele_probs = get_map_alleles(imputed_allele_probs)