diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py index 5fecbc0a39..a37492f739 100644 --- a/python/tests/beagle_numba.py +++ b/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)