Add assertions

python/tests/beagle_numba.py

@@ -826,7 +826,7 @@ def compute_individual_scores(
     :param numpy.ndarray allele_probs_1: Imputed allele probabilities for haplotype 1.
     :param numpy.ndarray alleles_2: Imputed alleles for haplotype 2.
     :param numpy.ndarray allele_probs_2: Imputed allele probabilities for haplotype 2.
-    :param int ref: REF allele.
+    :param int ref: Designated REF allele (ACGT encoding).
     :return: Genotype probabilities and allele dosages.
     :rtype: tuple(numpy.ndarray, numpy.ndarray)
     """
@@ -881,6 +881,7 @@ def compute_estimated_allelic_r_squared(gt_probs):
     _MIN_R2_DEN = 1e-8
     n = len(gt_probs)  # Number of individuals
     assert n > 0, "There must be at least one individual."
+    assert gt_probs.shape[1] == 3, "Three genotypes are considered."
     z = np.argmax(gt_probs, axis=1)  # Most likely imputed genotypes
     u = gt_probs[:, 1] + 2 * gt_probs[:, 2]  # E[X | y_i]
     w = gt_probs[:, 1] + 4 * gt_probs[:, 2]  # E[X^2 | y_i]
@@ -912,6 +913,7 @@ def compute_dosage_r_squared(gt_probs):
     _MIN_R2_DEN = 1e-8
     n = len(gt_probs)  # Number of individuals.
     assert n > 0, "There must be at least one individual."
+    assert gt_probs.shape[1] == 3, "Three genotypes are considered."
     u = gt_probs[:, 1] + 2 * gt_probs[:, 2]  # E[X | y_i].
     w = gt_probs[:, 1] + 4 * gt_probs[:, 2]  # E[X^2 | y_i].
     c = (1 / n) * np.sum(u) ** 2