Reorganize

python/tests/beagle_numba.py

@@ -577,6 +577,7 @@ def run_tsimpute(
     return (imputed_alleles, max_allele_probs)
 
 
+# Individual-level data.
 def compute_genotype_probs(alleles_1, allele_probs_1, alleles_2, allele_probs_2):
     """
     Compute genotype probabilities of diploid individuals at an imputed position.
@@ -615,6 +616,33 @@ def compute_genotype_probs(alleles_1, allele_probs_1, alleles_2, allele_probs_2)
     return gt_probs
 
 
+def compute_dosage_scores(alleles_1, allele_probs_1, alleles_2, allele_probs_2):
+    """
+    Compute dosage scores at ungenotyped positions of an individual.
+
+    In BEAGLE 4.1 output, DS: "Estimated ALT dose [P(RA) + P(AA)".
+
+    :param numpy.ndarray alleles_1: Imputed alleles for haplotype 1.
+    :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.
+    :return: Dosage scores at ungenotyped positions.
+    :rtype: numpy.ndarray(dtype=np.float64)
+    """
+    n = len(alleles_1)
+    assert len(alleles_2) == n, "Lengths of alleles differ."
+    assert (
+        len(allele_probs_1) == n
+    ), "Lengths of alleles and allele probabilities differ."
+    assert (
+        len(allele_probs_2) == n
+    ), "Lengths of alleles and allele probabilities differ."
+    dosage_scores = np.zeros(n, dtype=np.float64)
+    # TODO: Figure out how BEAGLE does this.
+    return dosage_scores
+
+
+# Site-level statistics.
 def compute_estimated_allelic_r_squared(gt_probs):
     """
     Compute the estimated allelic R^2 for an imputed position.
@@ -658,33 +686,6 @@ def compute_dosage_r_squared():
     pass
 
 
-def compute_dosage_scores(
-    imputed_alleles_1, allele_probs_1, imputed_alleles_2, allele_probs_2
-):
-    """
-    Compute dosage scores at ungenotyped positions of an individual.
-
-    In BEAGLE 4.1 output, DS: "Estimated ALT dose [P(RA) + P(AA)".
-
-    :param numpy.ndarray imputed_alleles_1: Imputed alleles for haplotype 1.
-    :param numpy.ndarray allele_probs_1: Imputed allele probabilities for haplotype 1.
-    :param numpy.ndarray imputed_alleles_2: Imputed alleles for haplotype 2.
-    :param numpy.ndarray allele_probs_2: Imputed allele probabilities for haplotype 2.
-    :return: Dosage scores at ungenotyped positions.
-    :rtype: numpy.ndarray(dtype=np.float64)
-    """
-    assert len(imputed_alleles_1) == len(
-        imputed_alleles_2
-    ), "Lengths of imputed alleles differ."
-    assert len(allele_probs_1) == len(
-        allele_probs_2
-    ), "Lengths of allele probabilities differ."
-    x = len(imputed_alleles_1)
-    dosage_scores = np.zeros(x, dtype=np.float64)
-    # TODO: Figure out how BEAGLE does this.
-    return dosage_scores
-
-
 def compute_alt_allele_frequencies():
     """
     In BEAGLE 4.1, AF: "Estimated ALT Allele Frequencies".