WIP: Allelic R-squared

szhan · Feb 25, 2024 · 9a3132b · 9a3132b
1 parent 1b974b8
commit 9a3132b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py
@@ -620,9 +620,10 @@ def compute_estimated_allelic_r_squared(
     ), "Lengths of alleles and allele probabilities differ."
     # Genotype probabilities for ith sample out of n samples.
     # Dossges: P(RR) = 0, P(RA or AR) = 1, P(AA) = 2.
+    # TODO: Refactor so that genotyped are the input.
     y = np.zeros((n, 3), dtype=np.float64)
-    u = y[:, 1] + 2 * y[:, 2]  # E[X | y_i].
     z = np.argmax(y, axis=1)  # Most likely imputed genotype of ith sample.
+    u = y[:, 1] + 2 * y[:, 2]  # E[X | y_i].
     w = y[:, 1] + 4 * y[:, 2]  # E[X^2 | y_i].
     term_1 = (np.sum(z * u) - (1 / n) * (np.sum(u) * np.sum(z))) ** 2
     term_2 = np.sum(w) - (1 / n) * np.sum(u) ** 2