Expose precision

python/tests/beagle_numba.py

@@ -369,7 +369,7 @@ def run_beagle(ref_h, query_h, pos, miscall_rate=1e-4, ne=1e6, genetic_map=None)
     return (imputed_alleles, max_allele_prob)
 
 
-def run_tsimpute(ref_ts, query_h, pos, mu, rho, genetic_map=None):
+def run_tsimpute(ref_ts, query_h, pos, mu, rho, precision=22, genetic_map=None):
     """
     Run the BEAGLE 4.1 algorithm except that the forward and backward probability
     matrices are computed from a tree sequence.
@@ -381,6 +381,7 @@ def run_tsimpute(ref_ts, query_h, pos, mu, rho, genetic_map=None):
     :param numpy.ndarray pos: Site positions of all the markers.
     :param numpy.ndarray mu: Mutation rate.
     :param numpy.ndarray rho: Recombination rate.
+    :param int precision: Precision when running LS HMM (default = 22).
     :param msprime.RateMap genetic_map: Genetic map.
     :return: Imputed alleles and their associated probabilities.
     :rtype: tuple(numpy.ndarray, numpy.ndarray)
@@ -410,7 +411,11 @@ def run_tsimpute(ref_ts, query_h, pos, mu, rho, genetic_map=None):
     fm = _tskit.CompressedMatrix(ref_ts_genotyped._ll_tree_sequence)
     bm = _tskit.CompressedMatrix(ref_ts_genotyped._ll_tree_sequence)
     ls_hmm = _tskit.LsHmm(
-        ref_ts_genotyped._ll_tree_sequence, mu, rho, acgt_alleles=True
+        ref_ts_genotyped._ll_tree_sequence,
+        mu,
+        rho,
+        acgt_alleles=True,
+        precision=precision,
     )
     ls_hmm.forward_matrix(query_h_genotyped.T, fm)
     ls_hmm.backward_matrix(query_h_genotyped.T, fm.normalisation_factor, bm)