Rearrange

python/tests/beagle_numba.py

@@ -662,13 +662,13 @@ def run_interpolation_beagle(
     # Get genetic map positions of of genotyped and ungenotyped markers.
     cm_typed = convert_to_genetic_map_positions(pos_typed, genetic_map=genetic_map)
     cm_untyped = convert_to_genetic_map_positions(pos_untyped, genetic_map=genetic_map)
-    # Subset haplotypes.
-    ref_h_typed = ref_h[idx_typed, :]  # Genotyped positions.
-    ref_h_untyped = ref_h[idx_untyped, :]  # Ungenotyped positions.
-    query_h_typed = query_h[idx_typed]  # Genotyped positions.
     # Get HMM probabilities at genotyped positions.
     trans_probs = get_transition_probs(cm_typed, h=h, ne=ne)
     mismatch_probs = get_mismatch_probs(len(pos_typed), error_rate=error_rate)
+    # Subset haplotypes.
+    ref_h_typed = ref_h[idx_typed, :]
+    ref_h_untyped = ref_h[idx_untyped, :]
+    query_h_typed = query_h[idx_typed]
     # Compute matrices at genotyped positions.
     fwd_mat = compute_forward_matrix(
         ref_h_typed,
@@ -750,7 +750,7 @@ def run_tsimpute(
         "Check the reference and query haplotypes use the same allele encoding.",
         stacklevel=1,
     )
-    h = ref_ts.num_samples
+    h = ref_ts.num_samples  # Number of reference haplotypes.
     # Separate indices of genotyped and ungenotyped positions.
     idx_typed = np.where(query_h != tskit.MISSING_DATA)[0]
     idx_untyped = np.where(query_h == tskit.MISSING_DATA)[0]