aegis_sim.submodels.genetics.ploider

View Source

 1import numpy as np
 2from numba import njit, prange
 3
 4
 5@njit(parallel=True)
 6def _diploid_to_haploid_numba(c0, c1, dominance_factor):
 7    """Parallel numba kernel for diploid-to-haploid conversion.
 8
 9    Operates on uint8 views of bool arrays to avoid numba's bool limitations.
10    Returns float32 output: 1.0 for homozygous true, 0.0 for homozygous false,
11    dominance_factor for heterozygous.
12    """
13    out = np.empty(c0.shape, dtype=np.float32)
14    n0, n1, n2 = c0.shape
15    for i in prange(n0):
16        for j in range(n1):
17            for k in range(n2):
18                a = c0[i, j, k]
19                b = c1[i, j, k]
20                if a == b:
21                    out[i, j, k] = np.float32(a)
22                else:
23                    out[i, j, k] = dominance_factor
24    return out
25
26
27class Ploider:
28    """ """
29
30    def init(self, REPRODUCTION_MODE, DOMINANCE_FACTOR, PLOIDY):
31        self.REPRODUCTION_MODE = REPRODUCTION_MODE
32        self.DOMINANCE_FACTOR = DOMINANCE_FACTOR
33        self.y = PLOIDY
34
35        if REPRODUCTION_MODE == "sexual":
36            assert PLOIDY == 2, f"If reproduction is sexual, ploidy can only be 2, not {PLOIDY}."
37
38    def diploid_to_haploid(self, loci):
39        """Merge two arrays encoding two chromatids into one array.
40
41        Arguments:
42            loci: A bool numpy array with shape (population size, ploidy, genome length, BITS_PER_LOCUS)
43
44        Returns:
45            A float numpy array with shape (population size, genome length, BITS_PER_LOCUS)
46        """
47
48        assert len(loci.shape) == 4, len(loci.shape)  # e.g. (45, 2, 250, 8)
49        assert loci.shape[1] == 2, loci.shape[1]  # ploidy
50
51        arr = _diploid_to_haploid_numba(
52            loci[:, 0].view(np.uint8),
53            loci[:, 1].view(np.uint8),
54            np.float32(self.DOMINANCE_FACTOR),
55        )
56
57        assert len(arr.shape) == 3, len(arr.shape)
58
59        return arr
60
61
62ploider = Ploider()

class Ploider: View Source

28class Ploider:
29    """ """
30
31    def init(self, REPRODUCTION_MODE, DOMINANCE_FACTOR, PLOIDY):
32        self.REPRODUCTION_MODE = REPRODUCTION_MODE
33        self.DOMINANCE_FACTOR = DOMINANCE_FACTOR
34        self.y = PLOIDY
35
36        if REPRODUCTION_MODE == "sexual":
37            assert PLOIDY == 2, f"If reproduction is sexual, ploidy can only be 2, not {PLOIDY}."
38
39    def diploid_to_haploid(self, loci):
40        """Merge two arrays encoding two chromatids into one array.
41
42        Arguments:
43            loci: A bool numpy array with shape (population size, ploidy, genome length, BITS_PER_LOCUS)
44
45        Returns:
46            A float numpy array with shape (population size, genome length, BITS_PER_LOCUS)
47        """
48
49        assert len(loci.shape) == 4, len(loci.shape)  # e.g. (45, 2, 250, 8)
50        assert loci.shape[1] == 2, loci.shape[1]  # ploidy
51
52        arr = _diploid_to_haploid_numba(
53            loci[:, 0].view(np.uint8),
54            loci[:, 1].view(np.uint8),
55            np.float32(self.DOMINANCE_FACTOR),
56        )
57
58        assert len(arr.shape) == 3, len(arr.shape)
59
60        return arr

def init(self, REPRODUCTION_MODE, DOMINANCE_FACTOR, PLOIDY): View Source

31    def init(self, REPRODUCTION_MODE, DOMINANCE_FACTOR, PLOIDY):
32        self.REPRODUCTION_MODE = REPRODUCTION_MODE
33        self.DOMINANCE_FACTOR = DOMINANCE_FACTOR
34        self.y = PLOIDY
35
36        if REPRODUCTION_MODE == "sexual":
37            assert PLOIDY == 2, f"If reproduction is sexual, ploidy can only be 2, not {PLOIDY}."

def diploid_to_haploid(self, loci): View Source

39    def diploid_to_haploid(self, loci):
40        """Merge two arrays encoding two chromatids into one array.
41
42        Arguments:
43            loci: A bool numpy array with shape (population size, ploidy, genome length, BITS_PER_LOCUS)
44
45        Returns:
46            A float numpy array with shape (population size, genome length, BITS_PER_LOCUS)
47        """
48
49        assert len(loci.shape) == 4, len(loci.shape)  # e.g. (45, 2, 250, 8)
50        assert loci.shape[1] == 2, loci.shape[1]  # ploidy
51
52        arr = _diploid_to_haploid_numba(
53            loci[:, 0].view(np.uint8),
54            loci[:, 1].view(np.uint8),
55            np.float32(self.DOMINANCE_FACTOR),
56        )
57
58        assert len(arr.shape) == 3, len(arr.shape)
59
60        return arr

Merge two arrays encoding two chromatids into one array.

Arguments: loci: A bool numpy array with shape (population size, ploidy, genome length, BITS_PER_LOCUS)

Returns: A float numpy array with shape (population size, genome length, BITS_PER_LOCUS)

ploider = <Ploider object>