aegis_sim.bioreactor

  1import numpy as np
  2import logging
  3
  4from aegis_sim import variables
  5from aegis_sim import submodels
  6from aegis_sim.constants import VALID_CAUSES_OF_DEATH
  7from aegis_sim.dataclasses.population import Population
  8from aegis_sim.recording import recordingmanager
  9from aegis_sim.parameterization import parametermanager
 10from aegis_sim.submodels.resources.resources import resources
 11
 12
 13class Bioreactor:
 14    def __init__(self, population: Population):
 15        self.eggs: Population = None
 16        self.population: Population = population
 17        self._starvation_steps: int = 0        # consecutive steps where N > resources
 18        self._starvation_multiplier: float = 1.0  # (1 - STARVATION_PENALTY) ** _starvation_steps
 19
 20    ##############
 21    # MAIN LOGIC #
 22    ##############
 23
 24    def run_step(self):
 25        """Perform one step of simulation."""
 26
 27        # If extinct (no living individuals nor eggs left), do nothing
 28        if len(self) == 0:
 29            logging.debug("Population went extinct.")
 30            recordingmanager.summaryrecorder.extinct = True
 31            return
 32
 33        # Scavenge resources and update the starvation multiplier.
 34        # If N > resources: starvation counter increments and multiplier compounds.
 35        # If resources >= N: counter resets to 0 and multiplier returns to 1.0.
 36        self._scavenge_resources()
 37
 38        # Mortality sources
 39        self.mortalities()
 40        resources.replenish()
 41
 42        recordingmanager.popsizerecorder.write_before_reproduction(self.population)
 43        self.growth()  # size increase
 44        self.reproduction()  # reproduction
 45        self.age()  # age increment and potentially death
 46        self.hatch()
 47        submodels.architect.envdrift.evolve(step=variables.steps)
 48
 49        # Record data
 50        recordingmanager.popsizerecorder.write_after_reproduction(self.population)
 51        recordingmanager.popsizerecorder.write_egg_num_after_reproduction(self.eggs)
 52        recordingmanager.envdriftmaprecorder.write(step=variables.steps)
 53        recordingmanager.flushrecorder.collect("additive_age_structure", self.population.ages)  # population census
 54        recordingmanager.picklerecorder.write(self.population)
 55        recordingmanager.featherrecorder.write(self.population)
 56        recordingmanager.ancestryrecorder.write(self.population)
 57        recordingmanager.guirecorder.record(self.population)
 58        recordingmanager.flushrecorder.flush()
 59        recordingmanager.popgenstatsrecorder.write(
 60            self.population.genomes, self.population.phenotypes.extract(ages=self.population.ages, trait_name="muta")
 61        )  # TODO defers calculation of mutation rates; hacky
 62        recordingmanager.summaryrecorder.record_memuse()
 63        recordingmanager.terecorder.record(self.population.ages, "alive")
 64        recordingmanager.checkpointrecorder.write(self.population, self.eggs)
 65
 66    ###############
 67    # STEP LOGIC #
 68    ###############
 69
 70    def mortalities(self):
 71        for source in parametermanager.parameters.MORTALITY_ORDER:
 72            if source == "intrinsic":
 73                self.mortality_intrinsic()
 74            elif source == "abiotic":
 75                self.mortality_abiotic()
 76            elif source == "infection":
 77                self.mortality_infection()
 78            elif source == "predation":
 79                self.mortality_predation()
 80            elif source == "starvation":
 81                self.mortality_starvation()
 82            else:
 83                raise ValueError(f"Invalid source of mortality '{source}'")
 84
 85    def mortality_intrinsic(self):
 86        probs_surv = self.population.phenotypes.extract(ages=self.population.ages, trait_name="surv")
 87        effective_surv = probs_surv * self._starvation_multiplier
 88        age_hazard = submodels.frailty.modify(hazard=1 - effective_surv, ages=self.population.ages)
 89        mask_kill = variables.rng.random(len(probs_surv)) < age_hazard
 90        self._kill(mask_kill=mask_kill, causeofdeath="intrinsic")
 91
 92    def mortality_abiotic(self):
 93        hazard = submodels.abiotic(variables.steps)
 94        age_hazard = submodels.frailty.modify(hazard=hazard, ages=self.population.ages)
 95        mask_kill = variables.rng.random(len(self.population)) < age_hazard
 96        self._kill(mask_kill=mask_kill, causeofdeath="abiotic")
 97
 98    def mortality_infection(self):
 99        submodels.infection(self.population)
100        # TODO add age hazard
101        mask_kill = self.population.infection == -1
102        self._kill(mask_kill=mask_kill, causeofdeath="infection")
103
104    def mortality_predation(self):
105        probs_kill = submodels.predation(len(self))
106        # TODO add age hazard
107        mask_kill = variables.rng.random(len(self)) < probs_kill
108        self._kill(mask_kill=mask_kill, causeofdeath="predation")
109
110    def mortality_starvation(self):
111        # Starvation now acts by scaling surv and repr phenotypes via _resource_ratio
112        # (set at the top of run_step before mortalities). No separate kill step needed.
113        # This method is kept so "starvation" remains a valid MORTALITY_ORDER entry.
114        pass
115
116    def reproduction(self):
117        """Generate offspring of reproducing individuals.
118        Initial is set to 0.
119        """
120
121        # Check if fertile
122        mask_fertile = (
123            self.population.ages >= parametermanager.parameters.MATURATION_AGE
124        )  # Check if mature; mature if survived MATURATION_AGE full cycles
125        if parametermanager.parameters.REPRODUCTION_ENDPOINT > 0:
126            mask_menopausal = (
127                self.population.ages >= parametermanager.parameters.REPRODUCTION_ENDPOINT
128            )  # Check if menopausal; menopausal when lived through REPRODUCTION_ENDPOINT full cycles
129            mask_fertile = (mask_fertile) & (~mask_menopausal)
130
131        if not any(mask_fertile):
132            return
133
134        probs_repr = (
135            self.population.phenotypes.extract(ages=self.population.ages, trait_name="repr", part=mask_fertile)
136            * self._starvation_multiplier
137        )
138
139        # Binomial calculation
140        n = parametermanager.parameters.MAX_OFFSPRING_NUMBER
141        p = probs_repr
142
143        assert np.all(p <= 1)
144        assert np.all(p >= 0)
145        num_repr = variables.rng.binomial(n=n, p=p)
146        mask_repr = num_repr > 0
147
148        if sum(num_repr) == 0:
149            return
150
151        # Indices of reproducing individuals
152        who = np.repeat(np.arange(len(self.population)), num_repr)
153
154        # Count ages at reproduction
155        ages_repr = self.population.ages[who]
156        recordingmanager.flushrecorder.collect("age_at_birth", ages_repr)
157
158        # Increase births statistics
159        self.population.births += num_repr
160
161        # Generate offspring genomes
162        parental_genomes = self.population.genomes.get(individuals=who)
163        parental_sexes = self.population.sexes[who]
164        parental_ancestry = self.population.ancestry[who] if self.population.ancestry is not None else None
165
166        muta_prob = self.population.phenotypes.extract(ages=self.population.ages, trait_name="muta", part=mask_repr)[
167            mask_repr
168        ]
169        muta_prob = np.repeat(muta_prob, num_repr[mask_repr])
170
171        offspring_genomes, offspring_ancestry = submodels.reproduction.generate_offspring_genomes(
172            genomes=parental_genomes,
173            muta_prob=muta_prob,
174            ages=ages_repr,
175            parental_sexes=parental_sexes,
176            ancestry=parental_ancestry,
177        )
178        offspring_sexes = submodels.sexsystem.get_sex(len(offspring_genomes))
179
180        # Randomize order of newly laid egg attributes ..
181        # .. because the order will affect their probability to be removed because of limited carrying capacity
182        order = np.arange(len(offspring_sexes))
183        variables.rng.shuffle(order)
184        offspring_genomes = offspring_genomes[order]
185        offspring_sexes = offspring_sexes[order]
186        if offspring_ancestry is not None:
187            offspring_ancestry = offspring_ancestry[order]
188
189        # Make eggs
190        eggs = Population.make_eggs(
191            offspring_genomes=offspring_genomes,
192            step=variables.steps,
193            offspring_sexes=offspring_sexes,
194            parental_generations=np.zeros(len(offspring_sexes)),  # TODO replace with working calculation
195            offspring_ancestry=offspring_ancestry,
196        )
197        if self.eggs is None:
198            self.eggs = eggs
199        else:
200            self.eggs += eggs
201
202        if parametermanager.parameters.CARRYING_CAPACITY_EGGS is not None and len(self.eggs) > parametermanager.parameters.CARRYING_CAPACITY_EGGS:
203            indices = np.arange(len(self.eggs))[-parametermanager.parameters.CARRYING_CAPACITY_EGGS :]
204            # TODO biased
205            self.eggs *= indices
206
207    def growth(self):
208        # TODO use already scavenged resources to determine growth
209        # max_growth_potential = self.population.phenotypes.extract(ages=self.population.ages, trait_name="grow")
210        # gathered_resources = submodels.resources.scavenge(max_growth_potential)
211        # self.population.sizes += gathered_resources
212        self.population.sizes += 1
213
214    def age(self):
215        """Increase age of all by one and kill those that surpass age limit.
216        Age denotes the number of full cycles that an individual survived and reproduced.
217        AGE_LIMIT is the maximum number of full cycles an individual can go through.
218        """
219        self.population.ages += 1
220        mask_kill = self.population.ages >= parametermanager.parameters.AGE_LIMIT
221        self._kill(mask_kill=mask_kill, causeofdeath="age_limit")
222
223    def hatch(self):
224        """Turn eggs into living individuals"""
225
226        # If nothing to hatch
227        if self.eggs is None or len(self.eggs) == 0:
228            return
229
230        # If REPRODUCTION_REGULATION is True, only reproduce until MAX_POPULATION_SIZE
231        if parametermanager.parameters.REPRODUCTION_REGULATION:
232            current_population_size = len(self.population)
233            remaining_capacity = resources.capacity - current_population_size
234            # If no remaining capacity, do not reproduce
235            if remaining_capacity < 1:
236                self.eggs = None
237                return
238            elif remaining_capacity < len(self.eggs):
239                indices = variables.rng.choice(len(self.eggs), size=int(remaining_capacity), replace=False)
240                self.eggs *= indices
241
242        # If something to hatch
243        if (
244            (
245                parametermanager.parameters.INCUBATION_PERIOD == -1 and len(self.population) == 0
246            )  # hatch when everyone dead
247            or (parametermanager.parameters.INCUBATION_PERIOD == 0)  # hatch immediately
248            or (
249                parametermanager.parameters.INCUBATION_PERIOD > 0
250                and variables.steps % parametermanager.parameters.INCUBATION_PERIOD == 0
251            )  # hatch with delay
252        ):
253
254            self.eggs.phenotypes = submodels.architect.__call__(self.eggs.genomes)
255            self.population += self.eggs
256            self.eggs = None
257
258    ################
259    # HELPER FUNCS #
260    ################
261
262    def _scavenge_resources(self):
263        """Scavenge resources and update the compounding starvation multiplier.
264
265        If N > available resources (deficit): starvation counter increments by 1
266        and multiplier = (1 - STARVATION_PENALTY) ** counter.
267
268        If resources >= N: counter resets to 0 and multiplier returns to 1.0.
269
270        The multiplier is applied to each individual's age-specific surv and repr
271        phenotypes in mortality_intrinsic() and reproduction().
272        """
273        n = len(self.population)
274        if n == 0:
275            self._starvation_steps = 0
276            self._starvation_multiplier = 1.0
277            return
278
279        in_deficit = n > resources.capacity
280
281        recordingmanager.resourcerecorder.write_before_scavenging()
282        resources.scavenge(np.ones(n))
283        recordingmanager.resourcerecorder.write_after_scavenging()
284
285        if in_deficit:
286            self._starvation_steps += 1
287        else:
288            self._starvation_steps = 0
289
290        penalty = parametermanager.parameters.STARVATION_PENALTY
291        self._starvation_multiplier = (1.0 - penalty) ** self._starvation_steps
292
293    def _kill(self, mask_kill, causeofdeath):
294        """Kill individuals and record their data."""
295
296        assert causeofdeath in VALID_CAUSES_OF_DEATH
297
298        # Skip if no one to kill
299        if not any(mask_kill):
300            return
301
302        # Count ages at death
303        # if causeofdeath != "age_limit":
304        ages_death = self.population.ages[mask_kill]
305        recordingmanager.flushrecorder.collect(f"age_at_{causeofdeath}", ages_death)
306        recordingmanager.terecorder.record(ages_death, "dead")
307
308        # Retain survivors
309        self.population *= ~mask_kill
310
311    def __len__(self):
312        """Return the number of living individuals and saved eggs."""
313        return len(self.population) + len(self.eggs) if self.eggs is not None else len(self.population)