Return marker_interval_start

python/tests/beagle_numba.py

@@ -52,6 +52,7 @@ def get_weights(genotyped_pos, imputed_pos):
     # Calculate weights for imputed markers.
     weights = np.zeros(x, dtype=np.float64)
     # Identify genotype markers k and k + 1 between imputed marker i
+    marker_interval_start = np.repeat(-2, x, dtype=np.int64)
     idx_m = m - 1
     for idx_x in np.arange(x - 1, -1, -1):
         assert (
@@ -72,9 +73,11 @@ def get_weights(genotyped_pos, imputed_pos):
             # Set min genetic distance to avoid division by zero.
             cm_mP1_m = cm_mP1_m if cm_mP1_m > 0.005 else 0.005
             weights[idx_x] = cm_mP1_x / cm_mP1_m
+        marker_interval_start[idx_x] = idx_m
     assert 0 <= np.min(weights), "Some weights are negative."
     assert np.max(weights) <= 1, "Some weights are greater than 1."
-    return weights
+    assert np.min(marker_interval_start) > -2, "Some marker indices are out of bounds."
+    return (weights, marker_interval_start)
 
 
 @njit
@@ -291,8 +294,8 @@ def interpolate_allele_probabilities(
     assert (x, h) == ref_h.shape
     genotyped_cm = convert_to_genetic_map_position(genotyped_pos)
     imputed_cm = convert_to_genetic_map_position(imputed_pos)
-    weights = get_weights(genotyped_pos, imputed_pos)
-    assert x == len(weights)
+    weights, marker_interval_start = get_weights(genotyped_pos, imputed_pos)
+    assert x == len(weights) == len(marker_interval_start)
     p = np.zeros((x, len(alleles)), dtype=np.float64)
     # Compute allele probabilities as per Equation 1 in BB2016.
     for a in alleles: