aegis_sim.submodels.genetics.ploider

 1import numpy as np
 2from numba import njit, prange
 3
 4
 5@njit(parallel=True)
 6def _diploid_to_haploid_numba(c0, c1, dominance_per_locus):
 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_per_locus[j] for heterozygous (j = locus index along axis 1).
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            h = dominance_per_locus[j]
18            for k in range(n2):
19                a = c0[i, j, k]
20                b = c1[i, j, k]
21                if a == b:
22                    out[i, j, k] = np.float32(a)
23                else:
24                    out[i, j, k] = h
25    return out
26
27
28class Ploider:
29    """ """
30
31    def init(self, REPRODUCTION_MODE, DOMINANCE_FACTOR, PLOIDY):
32        self.REPRODUCTION_MODE = REPRODUCTION_MODE
33        # Kept for backward-compatible config parsing; no longer used by the
34        # collapse kernel. See G_<trait>_dominance for per-trait control.
35        self.DOMINANCE_FACTOR = DOMINANCE_FACTOR
36        self.y = PLOIDY
37
38        if REPRODUCTION_MODE == "sexual":
39            assert PLOIDY == 2, f"If reproduction is sexual, ploidy can only be 2, not {PLOIDY}."
40
41    def diploid_to_haploid(self, loci, dominance_per_locus=None):
42        """Merge two arrays encoding two chromatids into one array.
43
44        Arguments:
45            loci: A bool numpy array with shape (population size, ploidy, genome length, BITS_PER_LOCUS)
46            dominance_per_locus: float32 array of shape (genome length,) giving h per locus.
47                If None, a uniform array of 0.5 (codominant) is used.
48
49        Returns:
50            A float numpy array with shape (population size, genome length, BITS_PER_LOCUS)
51        """
52
53        assert len(loci.shape) == 4, len(loci.shape)  # e.g. (45, 2, 250, 8)
54        assert loci.shape[1] == 2, loci.shape[1]  # ploidy
55
56        n_loci = loci.shape[2]
57        if dominance_per_locus is None:
58            # Backward-compatible fallback: callers that don't pass per-locus
59            # dominance (e.g. the modifying architecture, or unit tests) get a
60            # uniform array filled with self.DOMINANCE_FACTOR (the legacy global).
61            # The composite architecture passes a per-trait array that takes precedence.
62            fallback = getattr(self, "DOMINANCE_FACTOR", 0.5)
63            dominance_per_locus = np.full(n_loci, fallback, dtype=np.float32)
64        else:
65            assert dominance_per_locus.shape == (n_loci,), (
66                f"dominance_per_locus shape {dominance_per_locus.shape} != ({n_loci},)"
67            )
68            if dominance_per_locus.dtype != np.float32:
69                dominance_per_locus = dominance_per_locus.astype(np.float32)
70
71        arr = _diploid_to_haploid_numba(
72            loci[:, 0].view(np.uint8),
73            loci[:, 1].view(np.uint8),
74            dominance_per_locus,
75        )
76
77        assert len(arr.shape) == 3, len(arr.shape)
78
79        return arr
80
81
82ploider = Ploider()