Generalise to multiallelic sites

tskit-dev · Sep 12, 2023 · 65e7a1b · 65e7a1b
1 parent e30ee2a
commit 65e7a1b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/python/tests/beagle.py b/python/tests/beagle.py
@@ -498,7 +498,7 @@ def interpolate_allele_probabilities(
     imputed_cm = convert_to_genetic_map_position(imputed_pos)
     weights = get_weights(genotyped_pos, imputed_pos)
     assert x == len(weights)
-    p = np.zeros((x, 2), dtype=np.float64)
+    p = np.zeros((x, len(alleles)), dtype=np.float64)
     # Compute allele probabilities as per Equation 1 in BB2016.
     for a in alleles:
         for i in np.arange(x):

diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py
@@ -289,7 +289,7 @@ def interpolate_allele_probabilities(
     imputed_cm = convert_to_genetic_map_position(imputed_pos)
     weights = get_weights(genotyped_pos, imputed_pos)
     assert x == len(weights)
-    p = np.zeros((x, 2), dtype=np.float64)
+    p = np.zeros((x, len(alleles)), dtype=np.float64)
     # Compute allele probabilities as per Equation 1 in BB2016.
     for a in alleles:
         for i in np.arange(x):