aegis_sim.recording.ancestryrecorder

 1import pathlib
 2import numpy as np
 3import pandas as pd
 4
 5from aegis_sim import variables, parameterization
 6from aegis_sim.parameterization import parametermanager
 7from aegis_sim.constants import GENETIC_TRAITS
 8from aegis_sim.utilities.funcs import skip
 9from .recorder import Recorder
10
11
12class AncestryRecorder(Recorder):
13    """Records mean introgression fraction per locus at each snapshot step.
14
15    Output: snapshots/ancestry/{step}.csv
16    Rows: one per step. Columns: {trait}_{age} for each evolvable trait.
17    Values: mean fraction of introgressed alleles at that locus across all
18    living individuals (averaged over ploidy and population).
19    Only active when population.ancestry is not None (i.e. INTROGRESSION_SEEDS > 0).
20    """
21
22    def __init__(self, odir: pathlib.Path):
23        self.odir = odir / "snapshots" / "ancestry"
24        self.init_dir(self.odir)
25
26    def write(self, population):
27        if population.ancestry is None:
28            return
29        if skip("SNAPSHOT_RATE") or len(population) == 0:
30            return
31
32        step = variables.steps
33
34        # ancestry shape: (n, ploidy, n_loci_physical, bpl)
35        # mean over ploidy and bpl dimensions → (n, n_loci_physical)
36        locus_ancestry = population.ancestry.mean(axis=(1, 3))  # (n, n_loci_physical)
37        # mean over individuals → (n_loci_physical,)
38        mean_locus_physical = locus_ancestry.mean(axis=0)
39        # Un-permute from physical to logical (trait×age) order
40        from aegis_sim import submodels
41        perm = submodels.architect.architecture.locus_permutation
42        mean_locus = mean_locus_physical[perm]
43
44        cols = self._locus_columns()
45        df = pd.DataFrame([mean_locus], columns=cols)
46        df.insert(0, "step", step)
47        df.to_csv(self.odir / f"{step}.csv", index=False)
48
49    @staticmethod
50    def _locus_columns():
51        AGE_LIMIT = parametermanager.parameters.AGE_LIMIT
52        cols = []
53        for trait_name in GENETIC_TRAITS:
54            trait = parameterization.traits[trait_name]
55            if not trait.evolvable:
56                continue
57            if trait.agespecific:
58                cols.extend(f"{trait_name}_{age}" for age in range(AGE_LIMIT))
59            else:
60                cols.append(trait_name)
61        return cols