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