Add some doc

python/tests/beagle_numba.py

@@ -326,8 +326,18 @@ def run_beagle(ref_h, query_h, pos, miscall_rate=0.0001, ne=1e6):
 
 def run_tsimpute(ref_ts, query_h, pos, mu, rho):
     """
-    TODO: Document this function.
+    Run the BEAGLE algorithm except that the forward and backward probability
+    matrices are computed from a tree sequence.
+
     TODO: Put this function elsewhere.
+
+    :param numpy.ndarray ref_h: Reference haplotypes.
+    :param numpy.ndarray query_h: One query haplotype.
+    :param numpy.ndarray pos: Site positions of all the markers.
+    :param numpy.ndarray mu: Mutation rate.
+    :param numpy.ndarray rho: Recombination rate.
+    :return: Imputed alleles and their associated probabilities.
+    :rtype: tuple(numpy.ndarray, numpy.ndarray)
     """
     # Set indices of markers.
     genotyped_site_idx = np.where(query_h != -1)[0]