aegis_sim.recording.selectionrecorder

 1import logging
 2import pathlib
 3
 4from .recorder import Recorder
 5from aegis_sim import variables, submodels
 6from aegis_sim.parameterization import parametermanager
 7
 8
 9HEADER = "step,n_alive,n_alleles,n_carriers,allele_freq\n"
10
11
12class SelectionRecorder(Recorder):
13    """Track the allele frequency at the injection locus over time.
14
15    Active only when ALLELE_INJECTION_STEP > 0. Records start at step 1
16    (so the baseline before injection is visible) and continue through
17    the rest of the simulation.
18
19    Output: /selection/selection.csv with columns
20        step,n_alive,n_alleles,n_carriers,allele_freq
21    where n_alleles = 2 * n_alive (diploid count of alleles at the locus),
22    n_carriers = count of chromatids carrying ALLELE_INJECTION_ALLELE at
23    the injection locus, allele_freq = n_carriers / n_alleles.
24    """
25
26    def __init__(self, odir: pathlib.Path):
27        self.odir = odir / "selection"
28        self.init_odir()
29        self._file_path = self.odir / "selection.csv"
30        self._file = None
31        # Cached physical locus + bit, populated on first write.
32        self._physical_locus = None
33        self._bit_in_locus = None
34        self._tracked_allele = None
35
36    def _resolve_locus(self):
37        if self._physical_locus is not None:
38            return True
39        from aegis_sim import parameterization
40
41        trait_name = parametermanager.parameters.ALLELE_INJECTION_TRAIT
42        age = int(parametermanager.parameters.ALLELE_INJECTION_AGE)
43        bit_in_locus = int(parametermanager.parameters.ALLELE_INJECTION_BIT)
44        allele = bool(int(parametermanager.parameters.ALLELE_INJECTION_ALLELE))
45
46        trait = parameterization.traits.get(trait_name)
47        if trait is None or trait.length == 0:
48            return False
49        if trait.agespecific is True:
50            if not (0 <= age < trait.length):
51                return False
52            logical_locus = trait.start + age
53        else:
54            logical_locus = trait.start
55
56        self._physical_locus = int(submodels.architect.architecture.locus_permutation[logical_locus])
57        self._bit_in_locus = bit_in_locus
58        self._tracked_allele = allele
59        return True
60
61    def _ensure_open(self):
62        if self._file is not None:
63            return
64        mode = "a" if self._file_path.exists() and self._file_path.stat().st_size > 0 else "w"
65        self._file = open(self._file_path, mode)
66        if mode == "w":
67            self._file.write(HEADER)
68
69    def write(self, population):
70        if parametermanager.parameters.ALLELE_INJECTION_STEP <= 0:
71            return
72        if not self._resolve_locus():
73            return
74        n = len(population)
75        if n == 0:
76            return
77
78        bits = population.genomes.array[:, :, self._physical_locus, self._bit_in_locus]
79        # bits shape: (n, 2)
80        n_alleles = bits.size  # n * 2
81        n_carriers = int((bits == self._tracked_allele).sum())
82        freq = n_carriers / n_alleles if n_alleles else 0.0
83
84        self._ensure_open()
85        self._file.write(f"{variables.steps},{n},{n_alleles},{n_carriers},{freq:.6f}\n")
86        self._file.flush()