WIP: ImpData

szhan · Feb 25, 2024 · 75b78d1 · 75b78d1
1 parent 7608cb4
commit 75b78d1
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py
@@ -98,10 +98,23 @@ def __post_init__(self):
     def num_sites(self):
         return len(self.site_pos)
 
+    @property
+    def num_samples(self):
+        return self.alleles.shape[0]
+
     @property
     def num_individuals(self):
         return self.alleles.shape[0] / 2
 
+    def get_alleles_at_site(self, i):
+        idx_hap1 = np.arange(0, self.num_samples, 2)
+        idx_hap2 = np.arange(1, self.num_samples, 2)
+        alleles_1 = self.alleles[idx_hap1, i]
+        allele_probs_1 = self.allele_probs[idx_hap1, i]
+        alleles_2 = self.alleles[idx_hap2, i]
+        allele_probs_2 = self.allele_probs[idx_hap2, i]
+        return alleles_1, allele_probs_1, alleles_2, allele_probs_2
+
 
 def remap_alleles(a):
     """