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