aegis_sim.dataclasses.population

  1import numpy as np
  2import pickle
  3import pathlib
  4
  5from aegis_sim.dataclasses.genomes import Genomes
  6from aegis_sim.dataclasses.phenotypes import Phenotypes
  7from aegis_sim import submodels
  8
  9
 10class Population:
 11    """Population data
 12
 13    Contains demographic, genetic and phenotypic data of living individuals.
 14    """
 15
 16    attrs = (
 17        "genomes",
 18        "ages",
 19        "births",
 20        "birthdays",
 21        "generations",
 22        "phenotypes",
 23        "infection",
 24        "sizes",
 25        "sexes",
 26        "ancestry",
 27    )
 28
 29    def __init__(
 30        self,
 31        genomes: Genomes,
 32        ages,
 33        births,
 34        birthdays,
 35        phenotypes: Phenotypes,
 36        infection,
 37        sizes,
 38        sexes,
 39        generations=None,
 40        ancestry=None,
 41    ):
 42        self.genomes = genomes
 43        self.ages = ages
 44        self.births = births
 45        self.birthdays = birthdays
 46        self.phenotypes = phenotypes
 47        self.infection = infection
 48        self.sizes = sizes
 49        self.sexes = sexes
 50        self.generations = generations
 51        # ancestry: bool array, same shape as genomes.array; True = introgressed from source pop.
 52        # None when introgression tracking is disabled (standard runs).
 53        self.ancestry = ancestry
 54
 55        assert isinstance(phenotypes, Phenotypes)
 56
 57        if not (
 58            len(genomes)
 59            == len(ages)
 60            == len(births)
 61            == len(birthdays)
 62            == len(phenotypes)
 63            == len(infection)
 64            == len(sizes)
 65            == len(sexes)
 66            # == len(generations)
 67        ):
 68            raise ValueError("Population attributes must have equal length")
 69
 70    def __setstate__(self, state):
 71        # Backward compat: pickles saved before ancestry field was added
 72        self.__dict__.update(state)
 73        if "ancestry" not in self.__dict__:
 74            self.ancestry = None
 75
 76    def __len__(self):
 77        """Return the number of living individuals."""
 78        return len(self.genomes)
 79
 80    def __getitem__(self, index):
 81        """Return a subpopulation."""
 82        return Population(
 83            genomes=self.genomes.get(individuals=index),
 84            ages=self.ages[index],
 85            births=self.births[index],
 86            birthdays=self.birthdays[index],
 87            phenotypes=self.phenotypes.get(individuals=index),
 88            infection=self.infection[index],
 89            sizes=self.sizes[index],
 90            sexes=self.sexes[index],
 91            generations=self.generations[index] if self.generations is not None else None,
 92            ancestry=self.ancestry[index] if self.ancestry is not None else None,
 93        )
 94
 95    def __imul__(self, index):
 96        """Redefine itself as its own subpopulation."""
 97        for attr in self.attrs:
 98            if attr == "genomes":
 99                self.genomes.keep(individuals=index)
100            elif attr == "phenotypes":
101                self.phenotypes.keep(individuals=index)
102            elif attr == "generations":
103                self.generations = None
104            elif attr == "ancestry":
105                if self.ancestry is not None:
106                    self.ancestry = self.ancestry[index]
107            else:
108                setattr(self, attr, getattr(self, attr)[index])
109        return self
110
111    def __iadd__(self, population):
112        """Merge with another population."""
113
114        for attr in self.attrs:
115            if attr == "genomes":
116                self.genomes.add(population.genomes)
117            elif attr == "phenotypes":
118                assert isinstance(population.phenotypes, Phenotypes)
119                self.phenotypes.add(population.phenotypes)
120            elif attr == "generations":
121                self.generations = None
122            elif attr == "ancestry":
123                if self.ancestry is not None and population.ancestry is not None:
124                    self.ancestry = np.concatenate([self.ancestry, population.ancestry])
125                elif self.ancestry is not None or population.ancestry is not None:
126                    # one side has ancestry tracking, the other doesn't — treat missing as all-native
127                    a = self.ancestry if self.ancestry is not None else np.zeros(self.genomes.array.shape, dtype=np.bool_)
128                    b = population.ancestry if population.ancestry is not None else np.zeros(population.genomes.array.shape, dtype=np.bool_)
129                    self.ancestry = np.concatenate([a, b])
130            else:
131                val = np.concatenate([getattr(self, attr), getattr(population, attr)])
132                setattr(self, attr, val)
133        return self
134
135    # def shuffle(self):
136    #     order = np.random.arange(len(self))
137    #     np.random.shuffle(order)
138    #     self *= order
139
140    @staticmethod
141    def load_pickle_from(path: pathlib.Path):
142        assert path.exists(), f"pickle_path {path} does not exist"
143        with open(path, "rb") as file_:
144            return pickle.load(file_)
145
146    def save_pickle_to(self, path):
147        with open(path, "wb") as file_:
148            pickle.dump(self, file_)
149
150    @staticmethod
151    def initialize(n, AGE_LIMIT):
152        genomes = Genomes(submodels.architect.architecture.init_genome_array(n))
153        ages = np.random.randint(low=0, high=AGE_LIMIT, size=n, dtype=np.int32)
154        births = np.zeros(n, dtype=np.int32)
155        birthdays = np.zeros(n, dtype=np.int32)
156        # generations = np.zeros(n, dtype=np.int32)
157        generations = None
158
159        phenotypes = submodels.architect.__call__(genomes)
160        assert isinstance(phenotypes, Phenotypes)
161
162        infection = np.zeros(n, dtype=np.int32)
163        sizes = np.zeros(n, dtype=np.float32)
164        sexes = submodels.sexsystem.get_sex(n)
165        return Population(
166            genomes=genomes,
167            ages=ages,
168            births=births,
169            birthdays=birthdays,
170            generations=generations,
171            phenotypes=phenotypes,
172            infection=infection,
173            sizes=sizes,
174            sexes=sexes,
175            ancestry=None,
176        )
177
178    @staticmethod
179    def make_eggs(offspring_genomes: Genomes, step, offspring_sexes, parental_generations, offspring_ancestry=None):
180        n = len(offspring_genomes)
181        eggs = Population(
182            genomes=offspring_genomes,
183            ages=np.zeros(n, dtype=np.int32),
184            births=np.zeros(n, dtype=np.int32),
185            birthdays=np.zeros(n, dtype=np.int32) + step,
186            generations=None,
187            phenotypes=Phenotypes.init_phenotype_array(n),
188            infection=np.zeros(n, dtype=np.int32),
189            sizes=np.zeros(n, dtype=np.float32),
190            sexes=offspring_sexes,
191            ancestry=offspring_ancestry,
192        )
193        return eggs