Update docstring

python/tests/beagle_numba.py

@@ -433,7 +433,7 @@ def run_interpolation_beagle(
     ref_h_typed = ref_h[typed_idx, :]
     ref_h_untyped = ref_h[untyped_idx, :]
     query_h_typed = query_h[typed_idx]
-    # Set rates at genotyped markers.
+    # Get rates at genotyped markers.
     h = ref_h.shape[1]
     rho = get_recombination_rates(typed_cm, h=h, ne=ne)
     mu = get_mutation_rates(typed_pos, error_rate=error_rate)
@@ -450,7 +450,6 @@ def run_interpolation_beagle(
         typed_cm=typed_cm,
         untyped_cm=untyped_cm,
     )
-    # Get MAP alleles at ungenotyped markers.
     imputed_alleles, max_allele_probs = get_map_alleles(int_allele_probs)
     return (imputed_alleles, max_allele_probs)
 
@@ -468,18 +467,25 @@ def run_tsimpute(
     Perform interpolation-style imputation, except that the forward and backward
     probability matrices are computed from a tree sequence.
 
+    Reference haplotypes and query haplotype are of size (m + x, h) and (m + x).
+
+    The physical positions of all the markers are an array of size (m + x).
+
+    This produces imputed alleles and their interpolated probabilities
+    at the ungenotyped markers of the query haplotype.
+
     TODO: Put this function elsewhere.
     TODO: Set default precision. What should it be?
     TODO: Allow for imputation on user-specified genomic interval.
+    TODO: Handle `acgt_alleles` properly.
 
     :param numpy.ndarray ref_ts: Tree sequence with reference haplotypes.
     :param numpy.ndarray query_h: One query haplotype.
     :param numpy.ndarray pos: Physical positions of all the markers (bp).
     :param numpy.ndarray rho: Per-site recombination rate.
     :param numpy.ndarray mu: Per-site mutation rate.
-    :param int precision: Precision when running LS HMM (default = 22).
-    :param numpy.ndarray genetic_map_pos: Physical positions of genetic map (bp).
-    :param numpy.ndarray genetic_map_cm: Genetic map positions of genetic map (cM).
+    :param int precision: Precision when running Li & Stephens HMM.
+    :param GeneticMap genetic_map: Genetic map.
     :return: Imputed alleles and their interpolated probabilities.
     :rtype: tuple(numpy.ndarray, numpy.ndarray)
     """