From 34c3de2d2d1f8736b3c5e087e59a2ee4b5300ec9 Mon Sep 17 00:00:00 2001 From: Shing Zhan Date: Wed, 6 Mar 2024 07:17:44 +0000 Subject: [PATCH] Add assertions --- python/tests/beagle_numba.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py index 4c13f41a38..1f66b8d7bf 100644 --- a/python/tests/beagle_numba.py +++ b/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