Update comments in get_transition_probs

python/tests/beagle_numba.py

@@ -280,32 +280,32 @@ def get_mismatch_probs(pos, error_rate):
 
 def get_transition_probs(cm, h, ne):
     """
-    Compute transition probabilities at genotyped positions.
+    Compute probabilities of transitioning to a random state at genotyped positions.
 
     Note that in BEAGLE 4.1, the default value of `ne` is set to 1e6,
     whereas in BEAGLE 5.4, the default value of `ne` is set to 1e5.
-    In BEAGLE, this default value was optimized empirically on
+    In BEAGLE 4.1/5.4, this default value was optimized empirically on
     datasets from large outbred human populations. Also, in IMPUTE5,
     the default value of `ne` is set to 1e6.
 
     If `h` is small and `ne` is large, the transition probabilities are ~1.
     In such cases, it may be desirable to set `ne` to a small value
     to avoid switching between haplotypes all the time.
 
-    When using `_tskit.LsHmm`, this is `rho`.
+    This corresponds to `rho` in `_tskit.LsHmm`.
 
     :param numpy.ndarray cm: Genetic map positions (cM).
     :param int h: Number of reference haplotypes.
     :param float ne: Effective population size.
     :return: Transition probabilities.
     :rtype: numpy.ndarray
     """
-    # Recombination rate of first site is always 0.
+    # E(number of crossover events) at first site is always 0.
     r = np.zeros(len(cm), dtype=np.float64)
     r[1:] = np.diff(cm)
     c = -0.04 * (ne / h)
-    rho = -1 * np.expm1(c * r)
-    return rho
+    trans_probs = -1 * np.expm1(c * r)
+    return trans_probs
 
 
 @njit