From 75b78d18cdff23612e448ca32fd7cca4642781ab Mon Sep 17 00:00:00 2001 From: Shing Zhan Date: Sun, 25 Feb 2024 20:18:32 +0000 Subject: [PATCH] WIP: ImpData --- python/tests/beagle_numba.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py index 237a1ca5ca..802ff23200 100644 --- 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): """