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        # Selection-coefficient experiment: forced allele introduction at a specific step.
 39        if variables.steps == parametermanager.parameters.ALLELE_INJECTION_STEP:
 40            self._inject_allele()
 41
 42        # Selection-coefficient experiment: log allele frequency at the introduction locus.
 43        recordingmanager.selectionrecorder.write(self.population)
 44
 45        # Mortality sources
 46        self.mortalities()
 47        resources.replenish()
 48
 49        # Spatial lattice: migrate surviving individuals before reproduction.
 50        # Dead individuals' cells have been vacated by _kill via resync; this
 51        # step gives the survivors a chance to move into new cells. No-op
 52        # when LATTICE_MODE is False.
 53        if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
 54            submodels.lattice.migrate(
 55                positions=self.population.positions,
 56                migration_rate=parametermanager.parameters.MIGRATION_RATE,
 57                migration_long_rate=parametermanager.parameters.MIGRATION_LONG_RATE,
 58            )
 59
 60        recordingmanager.popsizerecorder.write_before_reproduction(self.population)
 61        self.growth()  # size increase
 62        self.reproduction()  # reproduction
 63        self.age()  # age increment and potentially death
 64        self.hatch()
 65        submodels.architect.envdrift.evolve(step=variables.steps)
 66
 67        # Record data
 68        recordingmanager.popsizerecorder.write_after_reproduction(self.population)
 69        recordingmanager.popsizerecorder.write_egg_num_after_reproduction(self.eggs)
 70        recordingmanager.envdriftmaprecorder.write(step=variables.steps)
 71        recordingmanager.flushrecorder.collect("additive_age_structure", self.population.ages)  # population census
 72        recordingmanager.picklerecorder.write(self.population)
 73        recordingmanager.featherrecorder.write(self.population)
 74        recordingmanager.ancestryrecorder.write(self.population)
 75        recordingmanager.fastarecorder.write(self.population)
 76        recordingmanager.vcfrecorder.write(self.population)
 77        recordingmanager.gvcfrecorder.write(self.population)
 78        recordingmanager.latticerecorder.write(self.population)
 79        recordingmanager.guirecorder.record(self.population)
 80        recordingmanager.flushrecorder.flush()
 81        recordingmanager.popgenstatsrecorder.write(
 82            self.population.genomes, self.population.phenotypes.extract(ages=self.population.ages, trait_name="muta")
 83        )  # TODO defers calculation of mutation rates; hacky
 84        recordingmanager.summaryrecorder.record_memuse()
 85        recordingmanager.terecorder.record(self.population.ages, "alive")
 86        recordingmanager.checkpointrecorder.write(self.population, self.eggs)
 87
 88    ###############
 89    # STEP LOGIC #
 90    ###############
 91
 92    def mortalities(self):
 93        for source in parametermanager.parameters.MORTALITY_ORDER:
 94            if source == "intrinsic":
 95                self.mortality_intrinsic()
 96            elif source == "abiotic":
 97                self.mortality_abiotic()
 98            elif source == "infection":
 99                self.mortality_infection()
100            elif source == "predation":
101                self.mortality_predation()
102            elif source == "starvation":
103                self.mortality_starvation()
104            else:
105                raise ValueError(f"Invalid source of mortality '{source}'")
106
107    def mortality_intrinsic(self):
108        probs_surv = self.population.phenotypes.extract(ages=self.population.ages, trait_name="surv")
109        effective_surv = probs_surv * self._starvation_multiplier
110        age_hazard = submodels.frailty.modify(hazard=1 - effective_surv, ages=self.population.ages)
111        mask_kill = variables.rng.random(len(probs_surv)) < age_hazard
112        self._kill(mask_kill=mask_kill, causeofdeath="intrinsic")
113
114    def mortality_abiotic(self):
115        hazard = submodels.abiotic(variables.steps)
116        age_hazard = submodels.frailty.modify(hazard=hazard, ages=self.population.ages)
117        mask_kill = variables.rng.random(len(self.population)) < age_hazard
118        self._kill(mask_kill=mask_kill, causeofdeath="abiotic")
119
120    def mortality_infection(self):
121        submodels.infection(self.population)
122        # TODO add age hazard
123        mask_kill = self.population.infection == -1
124        self._kill(mask_kill=mask_kill, causeofdeath="infection")
125
126    def mortality_predation(self):
127        probs_kill = submodels.predation(len(self))
128        # TODO add age hazard
129        mask_kill = variables.rng.random(len(self)) < probs_kill
130        self._kill(mask_kill=mask_kill, causeofdeath="predation")
131
132    def mortality_starvation(self):
133        # Starvation now acts by scaling surv and repr phenotypes via _resource_ratio
134        # (set at the top of run_step before mortalities). No separate kill step needed.
135        # This method is kept so "starvation" remains a valid MORTALITY_ORDER entry.
136        pass
137
138    def reproduction(self):
139        """Generate offspring of reproducing individuals.
140        Initial is set to 0.
141        """
142
143        # Check if fertile
144        mask_fertile = (
145            self.population.ages >= parametermanager.parameters.MATURATION_AGE
146        )  # Check if mature; mature if survived MATURATION_AGE full cycles
147        if parametermanager.parameters.REPRODUCTION_ENDPOINT > 0:
148            mask_menopausal = (
149                self.population.ages >= parametermanager.parameters.REPRODUCTION_ENDPOINT
150            )  # Check if menopausal; menopausal when lived through REPRODUCTION_ENDPOINT full cycles
151            mask_fertile = (mask_fertile) & (~mask_menopausal)
152
153        if not any(mask_fertile):
154            return
155
156        probs_repr = (
157            self.population.phenotypes.extract(ages=self.population.ages, trait_name="repr", part=mask_fertile)
158            * self._starvation_multiplier
159        )
160
161        # Binomial calculation
162        n = parametermanager.parameters.MAX_OFFSPRING_NUMBER
163        p = probs_repr
164
165        assert np.all(p <= 1)
166        assert np.all(p >= 0)
167        num_repr = variables.rng.binomial(n=n, p=p)
168        mask_repr = num_repr > 0
169
170        if sum(num_repr) == 0:
171            return
172
173        # Indices of reproducing individuals
174        who = np.repeat(np.arange(len(self.population)), num_repr)
175
176        # Count ages at reproduction
177        ages_repr = self.population.ages[who]
178        recordingmanager.flushrecorder.collect("age_at_birth", ages_repr)
179
180        # Increase births statistics
181        self.population.births += num_repr
182
183        # Generate offspring genomes
184        parental_genomes = self.population.genomes.get(individuals=who)
185        parental_sexes = self.population.sexes[who]
186        parental_ancestry = self.population.ancestry[who] if self.population.ancestry is not None else None
187
188        muta_prob = self.population.phenotypes.extract(ages=self.population.ages, trait_name="muta", part=mask_repr)[
189            mask_repr
190        ]
191        muta_prob = np.repeat(muta_prob, num_repr[mask_repr])
192
193        # When LATTICE_MODE + sexual, pass per-slot positions and the search
194        # radius to generate_offspring_genomes so matingmanager can do
195        # expanding-ring pairing. Otherwise the classical well-mixed pairing
196        # runs (parent_positions=None).
197        repr_parent_positions = None
198        repr_max_radius = 0
199        if (parametermanager.parameters.LATTICE_MODE
200                and self.population.positions is not None
201                and parametermanager.parameters.REPRODUCTION_MODE == "sexual"):
202            repr_parent_positions = self.population.positions[who]
203            repr_max_radius = int(parametermanager.parameters.MATING_MAX_SEARCH_RADIUS)
204
205        offspring_genomes, offspring_ancestry, mother_slots = submodels.reproduction.generate_offspring_genomes(
206            genomes=parental_genomes,
207            muta_prob=muta_prob,
208            ages=ages_repr,
209            parental_sexes=parental_sexes,
210            ancestry=parental_ancestry,
211            parent_positions=repr_parent_positions,
212            max_search_radius=repr_max_radius,
213        )
214        offspring_sexes = submodels.sexsystem.get_sex(len(offspring_genomes))
215
216        # Spatial lattice: place each offspring at a random empty cell adjacent
217        # to its mother. If no adjacent empty cell, birth fails (filtered out
218        # below). No-op when LATTICE_MODE is False.
219        offspring_positions = None
220        placed = None
221        if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
222            asexual = parametermanager.parameters.REPRODUCTION_MODE == "asexual"
223            # Determine each offspring's mother's position.
224            if asexual and len(offspring_genomes) == len(who):
225                parent_positions = self.population.positions[who]
226            elif (not asexual) and mother_slots is not None and len(mother_slots) == len(offspring_genomes):
227                # Sexual + lattice: mother_slots[i] is the slot index into
228                # parental_genomes; who[mother_slots[i]] is the population
229                # index of the mother of offspring i.
230                mother_population_idx = who[mother_slots]
231                parent_positions = self.population.positions[mother_population_idx]
232            else:
233                # No parent-position mapping available (e.g. sexual without
234                # lattice-aware pairing): fall back to whole-lattice random
235                # placement so the sim still progresses, with less spatial
236                # clustering.
237                parent_positions = None
238
239            offspring_positions = np.full((len(offspring_genomes), 2), -1, dtype=np.int32)
240            placed = np.zeros(len(offspring_genomes), dtype=bool)
241            for i in range(len(offspring_genomes)):
242                if parent_positions is not None:
243                    cur_q, cur_r = int(parent_positions[i, 0]), int(parent_positions[i, 1])
244                    target = submodels.lattice.random_empty_adjacent(cur_q, cur_r)
245                else:
246                    target = submodels.lattice.random_empty_anywhere()
247                if target is None:
248                    continue  # birth fails (no space)
249                submodels.lattice.claim(target[0], target[1])
250                offspring_positions[i] = target
251                placed[i] = True
252
253            if not placed.all():
254                offspring_genomes = offspring_genomes[placed]
255                offspring_sexes = offspring_sexes[placed]
256                if offspring_ancestry is not None:
257                    offspring_ancestry = offspring_ancestry[placed]
258                offspring_positions = offspring_positions[placed]
259                if mother_slots is not None:
260                    mother_slots = mother_slots[placed]
261
262        # Lineage tracking — clonal lineages only make biological sense for
263        # asexual reproduction, where each offspring has a single parent. With
264        # sexual reproduction each individual is genetically a mix of two
265        # parents, so any single-parent "lineage_id" colouring is misleading
266        # (it shows e.g. matrilineal mtDNA-style descent only, not actual
267        # genetic ancestry). For sexual sims, reconstruct ancestry post-hoc
268        # from the genome snapshots (FASTA / VCF / pickle) using standard
269        # phylogenetic tools instead.
270        offspring_lineage_id = None
271        offspring_parent_lineage_id = None
272        if (parametermanager.parameters.LINEAGE_TRACING
273                and self.population.lineage_id is not None
274                and parametermanager.parameters.REPRODUCTION_MODE == "asexual"):
275            # offspring[i] descends from parent at self.population[who[i]];
276            # filter `who` by `placed` if lattice placement dropped any births.
277            if parametermanager.parameters.LATTICE_MODE and offspring_positions is not None:
278                who_filtered = who[placed]
279                if len(offspring_genomes) == len(who_filtered):
280                    offspring_parent_lineage_id = self.population.lineage_id[who_filtered].astype(np.int64)
281                    offspring_lineage_id = variables.next_lineage_ids(len(offspring_genomes))
282            else:
283                if len(offspring_genomes) == len(who):
284                    offspring_parent_lineage_id = self.population.lineage_id[who].astype(np.int64)
285                    offspring_lineage_id = variables.next_lineage_ids(len(offspring_genomes))
286
287        # Randomize order of newly laid egg attributes ..
288        # .. because the order will affect their probability to be removed because of limited carrying capacity
289        order = np.arange(len(offspring_sexes))
290        variables.rng.shuffle(order)
291        offspring_genomes = offspring_genomes[order]
292        offspring_sexes = offspring_sexes[order]
293        if offspring_ancestry is not None:
294            offspring_ancestry = offspring_ancestry[order]
295        if offspring_lineage_id is not None:
296            offspring_lineage_id = offspring_lineage_id[order]
297            offspring_parent_lineage_id = offspring_parent_lineage_id[order]
298            recordingmanager.lineagerecorder.write_births(
299                parent_lineage_ids=offspring_parent_lineage_id,
300                child_lineage_ids=offspring_lineage_id,
301                step=variables.steps,
302            )
303        if offspring_positions is not None:
304            offspring_positions = offspring_positions[order]
305
306        # Make eggs
307        eggs = Population.make_eggs(
308            offspring_genomes=offspring_genomes,
309            step=variables.steps,
310            offspring_sexes=offspring_sexes,
311            parental_generations=np.zeros(len(offspring_sexes)),  # TODO replace with working calculation
312            offspring_ancestry=offspring_ancestry,
313            offspring_lineage_id=offspring_lineage_id,
314            offspring_parent_lineage_id=offspring_parent_lineage_id,
315            offspring_positions=offspring_positions,
316        )
317        if self.eggs is None:
318            self.eggs = eggs
319        else:
320            self.eggs += eggs
321
322        if parametermanager.parameters.CARRYING_CAPACITY_EGGS is not None and len(self.eggs) > parametermanager.parameters.CARRYING_CAPACITY_EGGS:
323            indices = np.arange(len(self.eggs))[-parametermanager.parameters.CARRYING_CAPACITY_EGGS :]
324            # TODO biased
325            self.eggs *= indices
326            # Truncated eggs lose their claim on the lattice cells too.
327            if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
328                submodels.lattice.resync_occupancy_from_positions(
329                    self.population.positions, self.eggs.positions
330                )
331
332    def growth(self):
333        # TODO use already scavenged resources to determine growth
334        # max_growth_potential = self.population.phenotypes.extract(ages=self.population.ages, trait_name="grow")
335        # gathered_resources = submodels.resources.scavenge(max_growth_potential)
336        # self.population.sizes += gathered_resources
337        self.population.sizes += 1
338
339    def age(self):
340        """Increase age of all by one and kill those that surpass age limit.
341        Age denotes the number of full cycles that an individual survived and reproduced.
342        AGE_LIMIT is the maximum number of full cycles an individual can go through.
343        """
344        self.population.ages += 1
345        mask_kill = self.population.ages >= parametermanager.parameters.AGE_LIMIT
346        self._kill(mask_kill=mask_kill, causeofdeath="age_limit")
347
348    def hatch(self):
349        """Turn eggs into living individuals"""
350
351        # If nothing to hatch
352        if self.eggs is None or len(self.eggs) == 0:
353            return
354
355        # If REPRODUCTION_REGULATION is True, only reproduce until MAX_POPULATION_SIZE
356        if parametermanager.parameters.REPRODUCTION_REGULATION:
357            current_population_size = len(self.population)
358            remaining_capacity = resources.capacity - current_population_size
359            # If no remaining capacity, do not reproduce
360            if remaining_capacity < 1:
361                self.eggs = None
362                return
363            elif remaining_capacity < len(self.eggs):
364                indices = variables.rng.choice(len(self.eggs), size=int(remaining_capacity), replace=False)
365                self.eggs *= indices
366
367        # If something to hatch
368        if (
369            (
370                parametermanager.parameters.INCUBATION_PERIOD == -1 and len(self.population) == 0
371            )  # hatch when everyone dead
372            or (parametermanager.parameters.INCUBATION_PERIOD == 0)  # hatch immediately
373            or (
374                parametermanager.parameters.INCUBATION_PERIOD > 0
375                and variables.steps % parametermanager.parameters.INCUBATION_PERIOD == 0
376            )  # hatch with delay
377        ):
378
379            # Lattice mode requires eggs to carry positions, assigned at reproduction
380            # time. If they don't, fail loudly rather than silently corrupt the
381            # lattice with (-1, -1) sentinel positions. (Offspring placement on
382            # the lattice is the next commit.)
383            if parametermanager.parameters.LATTICE_MODE and self.eggs.positions is None:
384                raise NotImplementedError(
385                    "LATTICE_MODE=True requires offspring placement on the lattice, "
386                    "which is not yet wired into reproduction. Use LATTICE_MODE=False "
387                    "for now, or wait for the lattice-reproduction commit."
388                )
389
390            self.eggs.phenotypes = submodels.architect.__call__(self.eggs.genomes)
391            self.population += self.eggs
392            self.eggs = None
393
394            # Sync lattice occupancy after population growth so migration sees
395            # all hatched individuals. No-op when LATTICE_MODE is False.
396            if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
397                submodels.lattice.resync_occupancy_from_positions(self.population.positions)
398
399    ################
400    # HELPER FUNCS #
401    ################
402
403    def _scavenge_resources(self):
404        """Scavenge resources and update the compounding starvation multiplier.
405
406        If N > available resources (deficit): starvation counter increments by 1
407        and multiplier = (1 - STARVATION_PENALTY) ** counter.
408
409        If resources >= N: counter resets to 0 and multiplier returns to 1.0.
410
411        The multiplier is applied to each individual's age-specific surv and repr
412        phenotypes in mortality_intrinsic() and reproduction().
413        """
414        n = len(self.population)
415        if n == 0:
416            self._starvation_steps = 0
417            self._starvation_multiplier = 1.0
418            return
419
420        in_deficit = n > resources.capacity
421
422        recordingmanager.resourcerecorder.write_before_scavenging()
423        resources.scavenge(np.ones(n))
424        recordingmanager.resourcerecorder.write_after_scavenging()
425
426        if in_deficit:
427            self._starvation_steps += 1
428        else:
429            self._starvation_steps = 0
430
431        penalty = parametermanager.parameters.STARVATION_PENALTY
432        self._starvation_multiplier = (1.0 - penalty) ** self._starvation_steps
433
434    def _kill(self, mask_kill, causeofdeath):
435        """Kill individuals and record their data."""
436
437        assert causeofdeath in VALID_CAUSES_OF_DEATH
438
439        # Skip if no one to kill
440        if not any(mask_kill):
441            return
442
443        # Count ages at death
444        # if causeofdeath != "age_limit":
445        ages_death = self.population.ages[mask_kill]
446        recordingmanager.flushrecorder.collect(f"age_at_{causeofdeath}", ages_death)
447        recordingmanager.terecorder.record(ages_death, "dead")
448
449        if self.population.lineage_id is not None:
450            recordingmanager.lineagerecorder.write_deaths(
451                lineage_ids=self.population.lineage_id[mask_kill],
452                causeofdeath=causeofdeath,
453                step=variables.steps,
454            )
455
456        # Retain survivors
457        self.population *= ~mask_kill
458
459        # Keep the lattice's occupancy grid consistent. Include incubating
460        # eggs' positions so they remain claimed during incubation periods.
461        # No-op when LATTICE_MODE is False.
462        if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
463            eggs_positions = self.eggs.positions if self.eggs is not None else None
464            submodels.lattice.resync_occupancy_from_positions(
465                self.population.positions, eggs_positions
466            )
467
468    def __len__(self):
469        """Return the number of living individuals and saved eggs."""
470        return len(self.population) + len(self.eggs) if self.eggs is not None else len(self.population)
471
472    def _inject_allele(self):
473        """Set ALLELE_INJECTION_ALLELE at the (TRAIT, AGE, BIT) locus on chromatid 0
474        of a ALLELE_INJECTION_FRACTION-sized random subset of the living population.
475        Recomputes phenotypes for the whole population so the new allele is expressed
476        immediately."""
477        from aegis_sim import parameterization
478
479        n = len(self.population)
480        if n == 0:
481            logging.warning("ALLELE_INJECTION_STEP reached but population is empty; skipping.")
482            return
483
484        trait_name = parametermanager.parameters.ALLELE_INJECTION_TRAIT
485        age = int(parametermanager.parameters.ALLELE_INJECTION_AGE)
486        bit_in_locus = int(parametermanager.parameters.ALLELE_INJECTION_BIT)
487        allele = bool(int(parametermanager.parameters.ALLELE_INJECTION_ALLELE))
488        fraction = float(parametermanager.parameters.ALLELE_INJECTION_FRACTION)
489
490        trait = parameterization.traits.get(trait_name)
491        if trait is None or trait.length == 0:
492            raise ValueError(
493                f"ALLELE_INJECTION_TRAIT={trait_name!r} is not a valid evolvable trait (length=0)."
494            )
495        if trait.agespecific is True:
496            if not (0 <= age < trait.length):
497                raise ValueError(f"ALLELE_INJECTION_AGE={age} out of range [0, {trait.length}) for trait {trait_name}.")
498            logical_locus = trait.start + age
499        else:
500            logical_locus = trait.start
501
502        bits_per_locus = self.population.genomes.array.shape[-1]
503        if not (0 <= bit_in_locus < bits_per_locus):
504            raise ValueError(f"ALLELE_INJECTION_BIT={bit_in_locus} out of range [0, {bits_per_locus}).")
505
506        physical_locus = int(submodels.architect.architecture.locus_permutation[logical_locus])
507
508        n_carriers = max(1, int(round(n * fraction)))
509        indices = variables.rng.choice(n, size=n_carriers, replace=False)
510        self.population.genomes.array[indices, 0, physical_locus, bit_in_locus] = allele
511
512        # Recompute phenotypes for the whole population so the new allele is expressed
513        # by the affected individuals' current age slot (cheap; same call as init).
514        self.population.phenotypes = submodels.architect(self.population.genomes)
515
516        logging.info(
517            "Mutation introduced at step %d: trait=%s age=%d bit=%d allele=%d in %d/%d individuals (chromatid 0, physical_locus=%d).",
518            variables.steps, trait_name, age, bit_in_locus, int(allele), n_carriers, n, physical_locus,
519        )
class Bioreactor:
 14class Bioreactor:
 15    def __init__(self, population: Population):
 16        self.eggs: Population = None
 17        self.population: Population = population
 18        self._starvation_steps: int = 0        # consecutive steps where N > resources
 19        self._starvation_multiplier: float = 1.0  # (1 - STARVATION_PENALTY) ** _starvation_steps
 20
 21    ##############
 22    # MAIN LOGIC #
 23    ##############
 24
 25    def run_step(self):
 26        """Perform one step of simulation."""
 27
 28        # If extinct (no living individuals nor eggs left), do nothing
 29        if len(self) == 0:
 30            logging.debug("Population went extinct.")
 31            recordingmanager.summaryrecorder.extinct = True
 32            return
 33
 34        # Scavenge resources and update the starvation multiplier.
 35        # If N > resources: starvation counter increments and multiplier compounds.
 36        # If resources >= N: counter resets to 0 and multiplier returns to 1.0.
 37        self._scavenge_resources()
 38
 39        # Selection-coefficient experiment: forced allele introduction at a specific step.
 40        if variables.steps == parametermanager.parameters.ALLELE_INJECTION_STEP:
 41            self._inject_allele()
 42
 43        # Selection-coefficient experiment: log allele frequency at the introduction locus.
 44        recordingmanager.selectionrecorder.write(self.population)
 45
 46        # Mortality sources
 47        self.mortalities()
 48        resources.replenish()
 49
 50        # Spatial lattice: migrate surviving individuals before reproduction.
 51        # Dead individuals' cells have been vacated by _kill via resync; this
 52        # step gives the survivors a chance to move into new cells. No-op
 53        # when LATTICE_MODE is False.
 54        if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
 55            submodels.lattice.migrate(
 56                positions=self.population.positions,
 57                migration_rate=parametermanager.parameters.MIGRATION_RATE,
 58                migration_long_rate=parametermanager.parameters.MIGRATION_LONG_RATE,
 59            )
 60
 61        recordingmanager.popsizerecorder.write_before_reproduction(self.population)
 62        self.growth()  # size increase
 63        self.reproduction()  # reproduction
 64        self.age()  # age increment and potentially death
 65        self.hatch()
 66        submodels.architect.envdrift.evolve(step=variables.steps)
 67
 68        # Record data
 69        recordingmanager.popsizerecorder.write_after_reproduction(self.population)
 70        recordingmanager.popsizerecorder.write_egg_num_after_reproduction(self.eggs)
 71        recordingmanager.envdriftmaprecorder.write(step=variables.steps)
 72        recordingmanager.flushrecorder.collect("additive_age_structure", self.population.ages)  # population census
 73        recordingmanager.picklerecorder.write(self.population)
 74        recordingmanager.featherrecorder.write(self.population)
 75        recordingmanager.ancestryrecorder.write(self.population)
 76        recordingmanager.fastarecorder.write(self.population)
 77        recordingmanager.vcfrecorder.write(self.population)
 78        recordingmanager.gvcfrecorder.write(self.population)
 79        recordingmanager.latticerecorder.write(self.population)
 80        recordingmanager.guirecorder.record(self.population)
 81        recordingmanager.flushrecorder.flush()
 82        recordingmanager.popgenstatsrecorder.write(
 83            self.population.genomes, self.population.phenotypes.extract(ages=self.population.ages, trait_name="muta")
 84        )  # TODO defers calculation of mutation rates; hacky
 85        recordingmanager.summaryrecorder.record_memuse()
 86        recordingmanager.terecorder.record(self.population.ages, "alive")
 87        recordingmanager.checkpointrecorder.write(self.population, self.eggs)
 88
 89    ###############
 90    # STEP LOGIC #
 91    ###############
 92
 93    def mortalities(self):
 94        for source in parametermanager.parameters.MORTALITY_ORDER:
 95            if source == "intrinsic":
 96                self.mortality_intrinsic()
 97            elif source == "abiotic":
 98                self.mortality_abiotic()
 99            elif source == "infection":
100                self.mortality_infection()
101            elif source == "predation":
102                self.mortality_predation()
103            elif source == "starvation":
104                self.mortality_starvation()
105            else:
106                raise ValueError(f"Invalid source of mortality '{source}'")
107
108    def mortality_intrinsic(self):
109        probs_surv = self.population.phenotypes.extract(ages=self.population.ages, trait_name="surv")
110        effective_surv = probs_surv * self._starvation_multiplier
111        age_hazard = submodels.frailty.modify(hazard=1 - effective_surv, ages=self.population.ages)
112        mask_kill = variables.rng.random(len(probs_surv)) < age_hazard
113        self._kill(mask_kill=mask_kill, causeofdeath="intrinsic")
114
115    def mortality_abiotic(self):
116        hazard = submodels.abiotic(variables.steps)
117        age_hazard = submodels.frailty.modify(hazard=hazard, ages=self.population.ages)
118        mask_kill = variables.rng.random(len(self.population)) < age_hazard
119        self._kill(mask_kill=mask_kill, causeofdeath="abiotic")
120
121    def mortality_infection(self):
122        submodels.infection(self.population)
123        # TODO add age hazard
124        mask_kill = self.population.infection == -1
125        self._kill(mask_kill=mask_kill, causeofdeath="infection")
126
127    def mortality_predation(self):
128        probs_kill = submodels.predation(len(self))
129        # TODO add age hazard
130        mask_kill = variables.rng.random(len(self)) < probs_kill
131        self._kill(mask_kill=mask_kill, causeofdeath="predation")
132
133    def mortality_starvation(self):
134        # Starvation now acts by scaling surv and repr phenotypes via _resource_ratio
135        # (set at the top of run_step before mortalities). No separate kill step needed.
136        # This method is kept so "starvation" remains a valid MORTALITY_ORDER entry.
137        pass
138
139    def reproduction(self):
140        """Generate offspring of reproducing individuals.
141        Initial is set to 0.
142        """
143
144        # Check if fertile
145        mask_fertile = (
146            self.population.ages >= parametermanager.parameters.MATURATION_AGE
147        )  # Check if mature; mature if survived MATURATION_AGE full cycles
148        if parametermanager.parameters.REPRODUCTION_ENDPOINT > 0:
149            mask_menopausal = (
150                self.population.ages >= parametermanager.parameters.REPRODUCTION_ENDPOINT
151            )  # Check if menopausal; menopausal when lived through REPRODUCTION_ENDPOINT full cycles
152            mask_fertile = (mask_fertile) & (~mask_menopausal)
153
154        if not any(mask_fertile):
155            return
156
157        probs_repr = (
158            self.population.phenotypes.extract(ages=self.population.ages, trait_name="repr", part=mask_fertile)
159            * self._starvation_multiplier
160        )
161
162        # Binomial calculation
163        n = parametermanager.parameters.MAX_OFFSPRING_NUMBER
164        p = probs_repr
165
166        assert np.all(p <= 1)
167        assert np.all(p >= 0)
168        num_repr = variables.rng.binomial(n=n, p=p)
169        mask_repr = num_repr > 0
170
171        if sum(num_repr) == 0:
172            return
173
174        # Indices of reproducing individuals
175        who = np.repeat(np.arange(len(self.population)), num_repr)
176
177        # Count ages at reproduction
178        ages_repr = self.population.ages[who]
179        recordingmanager.flushrecorder.collect("age_at_birth", ages_repr)
180
181        # Increase births statistics
182        self.population.births += num_repr
183
184        # Generate offspring genomes
185        parental_genomes = self.population.genomes.get(individuals=who)
186        parental_sexes = self.population.sexes[who]
187        parental_ancestry = self.population.ancestry[who] if self.population.ancestry is not None else None
188
189        muta_prob = self.population.phenotypes.extract(ages=self.population.ages, trait_name="muta", part=mask_repr)[
190            mask_repr
191        ]
192        muta_prob = np.repeat(muta_prob, num_repr[mask_repr])
193
194        # When LATTICE_MODE + sexual, pass per-slot positions and the search
195        # radius to generate_offspring_genomes so matingmanager can do
196        # expanding-ring pairing. Otherwise the classical well-mixed pairing
197        # runs (parent_positions=None).
198        repr_parent_positions = None
199        repr_max_radius = 0
200        if (parametermanager.parameters.LATTICE_MODE
201                and self.population.positions is not None
202                and parametermanager.parameters.REPRODUCTION_MODE == "sexual"):
203            repr_parent_positions = self.population.positions[who]
204            repr_max_radius = int(parametermanager.parameters.MATING_MAX_SEARCH_RADIUS)
205
206        offspring_genomes, offspring_ancestry, mother_slots = submodels.reproduction.generate_offspring_genomes(
207            genomes=parental_genomes,
208            muta_prob=muta_prob,
209            ages=ages_repr,
210            parental_sexes=parental_sexes,
211            ancestry=parental_ancestry,
212            parent_positions=repr_parent_positions,
213            max_search_radius=repr_max_radius,
214        )
215        offspring_sexes = submodels.sexsystem.get_sex(len(offspring_genomes))
216
217        # Spatial lattice: place each offspring at a random empty cell adjacent
218        # to its mother. If no adjacent empty cell, birth fails (filtered out
219        # below). No-op when LATTICE_MODE is False.
220        offspring_positions = None
221        placed = None
222        if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
223            asexual = parametermanager.parameters.REPRODUCTION_MODE == "asexual"
224            # Determine each offspring's mother's position.
225            if asexual and len(offspring_genomes) == len(who):
226                parent_positions = self.population.positions[who]
227            elif (not asexual) and mother_slots is not None and len(mother_slots) == len(offspring_genomes):
228                # Sexual + lattice: mother_slots[i] is the slot index into
229                # parental_genomes; who[mother_slots[i]] is the population
230                # index of the mother of offspring i.
231                mother_population_idx = who[mother_slots]
232                parent_positions = self.population.positions[mother_population_idx]
233            else:
234                # No parent-position mapping available (e.g. sexual without
235                # lattice-aware pairing): fall back to whole-lattice random
236                # placement so the sim still progresses, with less spatial
237                # clustering.
238                parent_positions = None
239
240            offspring_positions = np.full((len(offspring_genomes), 2), -1, dtype=np.int32)
241            placed = np.zeros(len(offspring_genomes), dtype=bool)
242            for i in range(len(offspring_genomes)):
243                if parent_positions is not None:
244                    cur_q, cur_r = int(parent_positions[i, 0]), int(parent_positions[i, 1])
245                    target = submodels.lattice.random_empty_adjacent(cur_q, cur_r)
246                else:
247                    target = submodels.lattice.random_empty_anywhere()
248                if target is None:
249                    continue  # birth fails (no space)
250                submodels.lattice.claim(target[0], target[1])
251                offspring_positions[i] = target
252                placed[i] = True
253
254            if not placed.all():
255                offspring_genomes = offspring_genomes[placed]
256                offspring_sexes = offspring_sexes[placed]
257                if offspring_ancestry is not None:
258                    offspring_ancestry = offspring_ancestry[placed]
259                offspring_positions = offspring_positions[placed]
260                if mother_slots is not None:
261                    mother_slots = mother_slots[placed]
262
263        # Lineage tracking — clonal lineages only make biological sense for
264        # asexual reproduction, where each offspring has a single parent. With
265        # sexual reproduction each individual is genetically a mix of two
266        # parents, so any single-parent "lineage_id" colouring is misleading
267        # (it shows e.g. matrilineal mtDNA-style descent only, not actual
268        # genetic ancestry). For sexual sims, reconstruct ancestry post-hoc
269        # from the genome snapshots (FASTA / VCF / pickle) using standard
270        # phylogenetic tools instead.
271        offspring_lineage_id = None
272        offspring_parent_lineage_id = None
273        if (parametermanager.parameters.LINEAGE_TRACING
274                and self.population.lineage_id is not None
275                and parametermanager.parameters.REPRODUCTION_MODE == "asexual"):
276            # offspring[i] descends from parent at self.population[who[i]];
277            # filter `who` by `placed` if lattice placement dropped any births.
278            if parametermanager.parameters.LATTICE_MODE and offspring_positions is not None:
279                who_filtered = who[placed]
280                if len(offspring_genomes) == len(who_filtered):
281                    offspring_parent_lineage_id = self.population.lineage_id[who_filtered].astype(np.int64)
282                    offspring_lineage_id = variables.next_lineage_ids(len(offspring_genomes))
283            else:
284                if len(offspring_genomes) == len(who):
285                    offspring_parent_lineage_id = self.population.lineage_id[who].astype(np.int64)
286                    offspring_lineage_id = variables.next_lineage_ids(len(offspring_genomes))
287
288        # Randomize order of newly laid egg attributes ..
289        # .. because the order will affect their probability to be removed because of limited carrying capacity
290        order = np.arange(len(offspring_sexes))
291        variables.rng.shuffle(order)
292        offspring_genomes = offspring_genomes[order]
293        offspring_sexes = offspring_sexes[order]
294        if offspring_ancestry is not None:
295            offspring_ancestry = offspring_ancestry[order]
296        if offspring_lineage_id is not None:
297            offspring_lineage_id = offspring_lineage_id[order]
298            offspring_parent_lineage_id = offspring_parent_lineage_id[order]
299            recordingmanager.lineagerecorder.write_births(
300                parent_lineage_ids=offspring_parent_lineage_id,
301                child_lineage_ids=offspring_lineage_id,
302                step=variables.steps,
303            )
304        if offspring_positions is not None:
305            offspring_positions = offspring_positions[order]
306
307        # Make eggs
308        eggs = Population.make_eggs(
309            offspring_genomes=offspring_genomes,
310            step=variables.steps,
311            offspring_sexes=offspring_sexes,
312            parental_generations=np.zeros(len(offspring_sexes)),  # TODO replace with working calculation
313            offspring_ancestry=offspring_ancestry,
314            offspring_lineage_id=offspring_lineage_id,
315            offspring_parent_lineage_id=offspring_parent_lineage_id,
316            offspring_positions=offspring_positions,
317        )
318        if self.eggs is None:
319            self.eggs = eggs
320        else:
321            self.eggs += eggs
322
323        if parametermanager.parameters.CARRYING_CAPACITY_EGGS is not None and len(self.eggs) > parametermanager.parameters.CARRYING_CAPACITY_EGGS:
324            indices = np.arange(len(self.eggs))[-parametermanager.parameters.CARRYING_CAPACITY_EGGS :]
325            # TODO biased
326            self.eggs *= indices
327            # Truncated eggs lose their claim on the lattice cells too.
328            if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
329                submodels.lattice.resync_occupancy_from_positions(
330                    self.population.positions, self.eggs.positions
331                )
332
333    def growth(self):
334        # TODO use already scavenged resources to determine growth
335        # max_growth_potential = self.population.phenotypes.extract(ages=self.population.ages, trait_name="grow")
336        # gathered_resources = submodels.resources.scavenge(max_growth_potential)
337        # self.population.sizes += gathered_resources
338        self.population.sizes += 1
339
340    def age(self):
341        """Increase age of all by one and kill those that surpass age limit.
342        Age denotes the number of full cycles that an individual survived and reproduced.
343        AGE_LIMIT is the maximum number of full cycles an individual can go through.
344        """
345        self.population.ages += 1
346        mask_kill = self.population.ages >= parametermanager.parameters.AGE_LIMIT
347        self._kill(mask_kill=mask_kill, causeofdeath="age_limit")
348
349    def hatch(self):
350        """Turn eggs into living individuals"""
351
352        # If nothing to hatch
353        if self.eggs is None or len(self.eggs) == 0:
354            return
355
356        # If REPRODUCTION_REGULATION is True, only reproduce until MAX_POPULATION_SIZE
357        if parametermanager.parameters.REPRODUCTION_REGULATION:
358            current_population_size = len(self.population)
359            remaining_capacity = resources.capacity - current_population_size
360            # If no remaining capacity, do not reproduce
361            if remaining_capacity < 1:
362                self.eggs = None
363                return
364            elif remaining_capacity < len(self.eggs):
365                indices = variables.rng.choice(len(self.eggs), size=int(remaining_capacity), replace=False)
366                self.eggs *= indices
367
368        # If something to hatch
369        if (
370            (
371                parametermanager.parameters.INCUBATION_PERIOD == -1 and len(self.population) == 0
372            )  # hatch when everyone dead
373            or (parametermanager.parameters.INCUBATION_PERIOD == 0)  # hatch immediately
374            or (
375                parametermanager.parameters.INCUBATION_PERIOD > 0
376                and variables.steps % parametermanager.parameters.INCUBATION_PERIOD == 0
377            )  # hatch with delay
378        ):
379
380            # Lattice mode requires eggs to carry positions, assigned at reproduction
381            # time. If they don't, fail loudly rather than silently corrupt the
382            # lattice with (-1, -1) sentinel positions. (Offspring placement on
383            # the lattice is the next commit.)
384            if parametermanager.parameters.LATTICE_MODE and self.eggs.positions is None:
385                raise NotImplementedError(
386                    "LATTICE_MODE=True requires offspring placement on the lattice, "
387                    "which is not yet wired into reproduction. Use LATTICE_MODE=False "
388                    "for now, or wait for the lattice-reproduction commit."
389                )
390
391            self.eggs.phenotypes = submodels.architect.__call__(self.eggs.genomes)
392            self.population += self.eggs
393            self.eggs = None
394
395            # Sync lattice occupancy after population growth so migration sees
396            # all hatched individuals. No-op when LATTICE_MODE is False.
397            if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
398                submodels.lattice.resync_occupancy_from_positions(self.population.positions)
399
400    ################
401    # HELPER FUNCS #
402    ################
403
404    def _scavenge_resources(self):
405        """Scavenge resources and update the compounding starvation multiplier.
406
407        If N > available resources (deficit): starvation counter increments by 1
408        and multiplier = (1 - STARVATION_PENALTY) ** counter.
409
410        If resources >= N: counter resets to 0 and multiplier returns to 1.0.
411
412        The multiplier is applied to each individual's age-specific surv and repr
413        phenotypes in mortality_intrinsic() and reproduction().
414        """
415        n = len(self.population)
416        if n == 0:
417            self._starvation_steps = 0
418            self._starvation_multiplier = 1.0
419            return
420
421        in_deficit = n > resources.capacity
422
423        recordingmanager.resourcerecorder.write_before_scavenging()
424        resources.scavenge(np.ones(n))
425        recordingmanager.resourcerecorder.write_after_scavenging()
426
427        if in_deficit:
428            self._starvation_steps += 1
429        else:
430            self._starvation_steps = 0
431
432        penalty = parametermanager.parameters.STARVATION_PENALTY
433        self._starvation_multiplier = (1.0 - penalty) ** self._starvation_steps
434
435    def _kill(self, mask_kill, causeofdeath):
436        """Kill individuals and record their data."""
437
438        assert causeofdeath in VALID_CAUSES_OF_DEATH
439
440        # Skip if no one to kill
441        if not any(mask_kill):
442            return
443
444        # Count ages at death
445        # if causeofdeath != "age_limit":
446        ages_death = self.population.ages[mask_kill]
447        recordingmanager.flushrecorder.collect(f"age_at_{causeofdeath}", ages_death)
448        recordingmanager.terecorder.record(ages_death, "dead")
449
450        if self.population.lineage_id is not None:
451            recordingmanager.lineagerecorder.write_deaths(
452                lineage_ids=self.population.lineage_id[mask_kill],
453                causeofdeath=causeofdeath,
454                step=variables.steps,
455            )
456
457        # Retain survivors
458        self.population *= ~mask_kill
459
460        # Keep the lattice's occupancy grid consistent. Include incubating
461        # eggs' positions so they remain claimed during incubation periods.
462        # No-op when LATTICE_MODE is False.
463        if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
464            eggs_positions = self.eggs.positions if self.eggs is not None else None
465            submodels.lattice.resync_occupancy_from_positions(
466                self.population.positions, eggs_positions
467            )
468
469    def __len__(self):
470        """Return the number of living individuals and saved eggs."""
471        return len(self.population) + len(self.eggs) if self.eggs is not None else len(self.population)
472
473    def _inject_allele(self):
474        """Set ALLELE_INJECTION_ALLELE at the (TRAIT, AGE, BIT) locus on chromatid 0
475        of a ALLELE_INJECTION_FRACTION-sized random subset of the living population.
476        Recomputes phenotypes for the whole population so the new allele is expressed
477        immediately."""
478        from aegis_sim import parameterization
479
480        n = len(self.population)
481        if n == 0:
482            logging.warning("ALLELE_INJECTION_STEP reached but population is empty; skipping.")
483            return
484
485        trait_name = parametermanager.parameters.ALLELE_INJECTION_TRAIT
486        age = int(parametermanager.parameters.ALLELE_INJECTION_AGE)
487        bit_in_locus = int(parametermanager.parameters.ALLELE_INJECTION_BIT)
488        allele = bool(int(parametermanager.parameters.ALLELE_INJECTION_ALLELE))
489        fraction = float(parametermanager.parameters.ALLELE_INJECTION_FRACTION)
490
491        trait = parameterization.traits.get(trait_name)
492        if trait is None or trait.length == 0:
493            raise ValueError(
494                f"ALLELE_INJECTION_TRAIT={trait_name!r} is not a valid evolvable trait (length=0)."
495            )
496        if trait.agespecific is True:
497            if not (0 <= age < trait.length):
498                raise ValueError(f"ALLELE_INJECTION_AGE={age} out of range [0, {trait.length}) for trait {trait_name}.")
499            logical_locus = trait.start + age
500        else:
501            logical_locus = trait.start
502
503        bits_per_locus = self.population.genomes.array.shape[-1]
504        if not (0 <= bit_in_locus < bits_per_locus):
505            raise ValueError(f"ALLELE_INJECTION_BIT={bit_in_locus} out of range [0, {bits_per_locus}).")
506
507        physical_locus = int(submodels.architect.architecture.locus_permutation[logical_locus])
508
509        n_carriers = max(1, int(round(n * fraction)))
510        indices = variables.rng.choice(n, size=n_carriers, replace=False)
511        self.population.genomes.array[indices, 0, physical_locus, bit_in_locus] = allele
512
513        # Recompute phenotypes for the whole population so the new allele is expressed
514        # by the affected individuals' current age slot (cheap; same call as init).
515        self.population.phenotypes = submodels.architect(self.population.genomes)
516
517        logging.info(
518            "Mutation introduced at step %d: trait=%s age=%d bit=%d allele=%d in %d/%d individuals (chromatid 0, physical_locus=%d).",
519            variables.steps, trait_name, age, bit_in_locus, int(allele), n_carriers, n, physical_locus,
520        )
Bioreactor(population: aegis_sim.dataclasses.population.Population)
15    def __init__(self, population: Population):
16        self.eggs: Population = None
17        self.population: Population = population
18        self._starvation_steps: int = 0        # consecutive steps where N > resources
19        self._starvation_multiplier: float = 1.0  # (1 - STARVATION_PENALTY) ** _starvation_steps
def run_step(self):
25    def run_step(self):
26        """Perform one step of simulation."""
27
28        # If extinct (no living individuals nor eggs left), do nothing
29        if len(self) == 0:
30            logging.debug("Population went extinct.")
31            recordingmanager.summaryrecorder.extinct = True
32            return
33
34        # Scavenge resources and update the starvation multiplier.
35        # If N > resources: starvation counter increments and multiplier compounds.
36        # If resources >= N: counter resets to 0 and multiplier returns to 1.0.
37        self._scavenge_resources()
38
39        # Selection-coefficient experiment: forced allele introduction at a specific step.
40        if variables.steps == parametermanager.parameters.ALLELE_INJECTION_STEP:
41            self._inject_allele()
42
43        # Selection-coefficient experiment: log allele frequency at the introduction locus.
44        recordingmanager.selectionrecorder.write(self.population)
45
46        # Mortality sources
47        self.mortalities()
48        resources.replenish()
49
50        # Spatial lattice: migrate surviving individuals before reproduction.
51        # Dead individuals' cells have been vacated by _kill via resync; this
52        # step gives the survivors a chance to move into new cells. No-op
53        # when LATTICE_MODE is False.
54        if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
55            submodels.lattice.migrate(
56                positions=self.population.positions,
57                migration_rate=parametermanager.parameters.MIGRATION_RATE,
58                migration_long_rate=parametermanager.parameters.MIGRATION_LONG_RATE,
59            )
60
61        recordingmanager.popsizerecorder.write_before_reproduction(self.population)
62        self.growth()  # size increase
63        self.reproduction()  # reproduction
64        self.age()  # age increment and potentially death
65        self.hatch()
66        submodels.architect.envdrift.evolve(step=variables.steps)
67
68        # Record data
69        recordingmanager.popsizerecorder.write_after_reproduction(self.population)
70        recordingmanager.popsizerecorder.write_egg_num_after_reproduction(self.eggs)
71        recordingmanager.envdriftmaprecorder.write(step=variables.steps)
72        recordingmanager.flushrecorder.collect("additive_age_structure", self.population.ages)  # population census
73        recordingmanager.picklerecorder.write(self.population)
74        recordingmanager.featherrecorder.write(self.population)
75        recordingmanager.ancestryrecorder.write(self.population)
76        recordingmanager.fastarecorder.write(self.population)
77        recordingmanager.vcfrecorder.write(self.population)
78        recordingmanager.gvcfrecorder.write(self.population)
79        recordingmanager.latticerecorder.write(self.population)
80        recordingmanager.guirecorder.record(self.population)
81        recordingmanager.flushrecorder.flush()
82        recordingmanager.popgenstatsrecorder.write(
83            self.population.genomes, self.population.phenotypes.extract(ages=self.population.ages, trait_name="muta")
84        )  # TODO defers calculation of mutation rates; hacky
85        recordingmanager.summaryrecorder.record_memuse()
86        recordingmanager.terecorder.record(self.population.ages, "alive")
87        recordingmanager.checkpointrecorder.write(self.population, self.eggs)

Perform one step of simulation.

def mortalities(self):
 93    def mortalities(self):
 94        for source in parametermanager.parameters.MORTALITY_ORDER:
 95            if source == "intrinsic":
 96                self.mortality_intrinsic()
 97            elif source == "abiotic":
 98                self.mortality_abiotic()
 99            elif source == "infection":
100                self.mortality_infection()
101            elif source == "predation":
102                self.mortality_predation()
103            elif source == "starvation":
104                self.mortality_starvation()
105            else:
106                raise ValueError(f"Invalid source of mortality '{source}'")
def mortality_intrinsic(self):
108    def mortality_intrinsic(self):
109        probs_surv = self.population.phenotypes.extract(ages=self.population.ages, trait_name="surv")
110        effective_surv = probs_surv * self._starvation_multiplier
111        age_hazard = submodels.frailty.modify(hazard=1 - effective_surv, ages=self.population.ages)
112        mask_kill = variables.rng.random(len(probs_surv)) < age_hazard
113        self._kill(mask_kill=mask_kill, causeofdeath="intrinsic")
def mortality_abiotic(self):
115    def mortality_abiotic(self):
116        hazard = submodels.abiotic(variables.steps)
117        age_hazard = submodels.frailty.modify(hazard=hazard, ages=self.population.ages)
118        mask_kill = variables.rng.random(len(self.population)) < age_hazard
119        self._kill(mask_kill=mask_kill, causeofdeath="abiotic")
def mortality_infection(self):
121    def mortality_infection(self):
122        submodels.infection(self.population)
123        # TODO add age hazard
124        mask_kill = self.population.infection == -1
125        self._kill(mask_kill=mask_kill, causeofdeath="infection")
def mortality_predation(self):
127    def mortality_predation(self):
128        probs_kill = submodels.predation(len(self))
129        # TODO add age hazard
130        mask_kill = variables.rng.random(len(self)) < probs_kill
131        self._kill(mask_kill=mask_kill, causeofdeath="predation")
def mortality_starvation(self):
133    def mortality_starvation(self):
134        # Starvation now acts by scaling surv and repr phenotypes via _resource_ratio
135        # (set at the top of run_step before mortalities). No separate kill step needed.
136        # This method is kept so "starvation" remains a valid MORTALITY_ORDER entry.
137        pass
def reproduction(self):
139    def reproduction(self):
140        """Generate offspring of reproducing individuals.
141        Initial is set to 0.
142        """
143
144        # Check if fertile
145        mask_fertile = (
146            self.population.ages >= parametermanager.parameters.MATURATION_AGE
147        )  # Check if mature; mature if survived MATURATION_AGE full cycles
148        if parametermanager.parameters.REPRODUCTION_ENDPOINT > 0:
149            mask_menopausal = (
150                self.population.ages >= parametermanager.parameters.REPRODUCTION_ENDPOINT
151            )  # Check if menopausal; menopausal when lived through REPRODUCTION_ENDPOINT full cycles
152            mask_fertile = (mask_fertile) & (~mask_menopausal)
153
154        if not any(mask_fertile):
155            return
156
157        probs_repr = (
158            self.population.phenotypes.extract(ages=self.population.ages, trait_name="repr", part=mask_fertile)
159            * self._starvation_multiplier
160        )
161
162        # Binomial calculation
163        n = parametermanager.parameters.MAX_OFFSPRING_NUMBER
164        p = probs_repr
165
166        assert np.all(p <= 1)
167        assert np.all(p >= 0)
168        num_repr = variables.rng.binomial(n=n, p=p)
169        mask_repr = num_repr > 0
170
171        if sum(num_repr) == 0:
172            return
173
174        # Indices of reproducing individuals
175        who = np.repeat(np.arange(len(self.population)), num_repr)
176
177        # Count ages at reproduction
178        ages_repr = self.population.ages[who]
179        recordingmanager.flushrecorder.collect("age_at_birth", ages_repr)
180
181        # Increase births statistics
182        self.population.births += num_repr
183
184        # Generate offspring genomes
185        parental_genomes = self.population.genomes.get(individuals=who)
186        parental_sexes = self.population.sexes[who]
187        parental_ancestry = self.population.ancestry[who] if self.population.ancestry is not None else None
188
189        muta_prob = self.population.phenotypes.extract(ages=self.population.ages, trait_name="muta", part=mask_repr)[
190            mask_repr
191        ]
192        muta_prob = np.repeat(muta_prob, num_repr[mask_repr])
193
194        # When LATTICE_MODE + sexual, pass per-slot positions and the search
195        # radius to generate_offspring_genomes so matingmanager can do
196        # expanding-ring pairing. Otherwise the classical well-mixed pairing
197        # runs (parent_positions=None).
198        repr_parent_positions = None
199        repr_max_radius = 0
200        if (parametermanager.parameters.LATTICE_MODE
201                and self.population.positions is not None
202                and parametermanager.parameters.REPRODUCTION_MODE == "sexual"):
203            repr_parent_positions = self.population.positions[who]
204            repr_max_radius = int(parametermanager.parameters.MATING_MAX_SEARCH_RADIUS)
205
206        offspring_genomes, offspring_ancestry, mother_slots = submodels.reproduction.generate_offspring_genomes(
207            genomes=parental_genomes,
208            muta_prob=muta_prob,
209            ages=ages_repr,
210            parental_sexes=parental_sexes,
211            ancestry=parental_ancestry,
212            parent_positions=repr_parent_positions,
213            max_search_radius=repr_max_radius,
214        )
215        offspring_sexes = submodels.sexsystem.get_sex(len(offspring_genomes))
216
217        # Spatial lattice: place each offspring at a random empty cell adjacent
218        # to its mother. If no adjacent empty cell, birth fails (filtered out
219        # below). No-op when LATTICE_MODE is False.
220        offspring_positions = None
221        placed = None
222        if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
223            asexual = parametermanager.parameters.REPRODUCTION_MODE == "asexual"
224            # Determine each offspring's mother's position.
225            if asexual and len(offspring_genomes) == len(who):
226                parent_positions = self.population.positions[who]
227            elif (not asexual) and mother_slots is not None and len(mother_slots) == len(offspring_genomes):
228                # Sexual + lattice: mother_slots[i] is the slot index into
229                # parental_genomes; who[mother_slots[i]] is the population
230                # index of the mother of offspring i.
231                mother_population_idx = who[mother_slots]
232                parent_positions = self.population.positions[mother_population_idx]
233            else:
234                # No parent-position mapping available (e.g. sexual without
235                # lattice-aware pairing): fall back to whole-lattice random
236                # placement so the sim still progresses, with less spatial
237                # clustering.
238                parent_positions = None
239
240            offspring_positions = np.full((len(offspring_genomes), 2), -1, dtype=np.int32)
241            placed = np.zeros(len(offspring_genomes), dtype=bool)
242            for i in range(len(offspring_genomes)):
243                if parent_positions is not None:
244                    cur_q, cur_r = int(parent_positions[i, 0]), int(parent_positions[i, 1])
245                    target = submodels.lattice.random_empty_adjacent(cur_q, cur_r)
246                else:
247                    target = submodels.lattice.random_empty_anywhere()
248                if target is None:
249                    continue  # birth fails (no space)
250                submodels.lattice.claim(target[0], target[1])
251                offspring_positions[i] = target
252                placed[i] = True
253
254            if not placed.all():
255                offspring_genomes = offspring_genomes[placed]
256                offspring_sexes = offspring_sexes[placed]
257                if offspring_ancestry is not None:
258                    offspring_ancestry = offspring_ancestry[placed]
259                offspring_positions = offspring_positions[placed]
260                if mother_slots is not None:
261                    mother_slots = mother_slots[placed]
262
263        # Lineage tracking — clonal lineages only make biological sense for
264        # asexual reproduction, where each offspring has a single parent. With
265        # sexual reproduction each individual is genetically a mix of two
266        # parents, so any single-parent "lineage_id" colouring is misleading
267        # (it shows e.g. matrilineal mtDNA-style descent only, not actual
268        # genetic ancestry). For sexual sims, reconstruct ancestry post-hoc
269        # from the genome snapshots (FASTA / VCF / pickle) using standard
270        # phylogenetic tools instead.
271        offspring_lineage_id = None
272        offspring_parent_lineage_id = None
273        if (parametermanager.parameters.LINEAGE_TRACING
274                and self.population.lineage_id is not None
275                and parametermanager.parameters.REPRODUCTION_MODE == "asexual"):
276            # offspring[i] descends from parent at self.population[who[i]];
277            # filter `who` by `placed` if lattice placement dropped any births.
278            if parametermanager.parameters.LATTICE_MODE and offspring_positions is not None:
279                who_filtered = who[placed]
280                if len(offspring_genomes) == len(who_filtered):
281                    offspring_parent_lineage_id = self.population.lineage_id[who_filtered].astype(np.int64)
282                    offspring_lineage_id = variables.next_lineage_ids(len(offspring_genomes))
283            else:
284                if len(offspring_genomes) == len(who):
285                    offspring_parent_lineage_id = self.population.lineage_id[who].astype(np.int64)
286                    offspring_lineage_id = variables.next_lineage_ids(len(offspring_genomes))
287
288        # Randomize order of newly laid egg attributes ..
289        # .. because the order will affect their probability to be removed because of limited carrying capacity
290        order = np.arange(len(offspring_sexes))
291        variables.rng.shuffle(order)
292        offspring_genomes = offspring_genomes[order]
293        offspring_sexes = offspring_sexes[order]
294        if offspring_ancestry is not None:
295            offspring_ancestry = offspring_ancestry[order]
296        if offspring_lineage_id is not None:
297            offspring_lineage_id = offspring_lineage_id[order]
298            offspring_parent_lineage_id = offspring_parent_lineage_id[order]
299            recordingmanager.lineagerecorder.write_births(
300                parent_lineage_ids=offspring_parent_lineage_id,
301                child_lineage_ids=offspring_lineage_id,
302                step=variables.steps,
303            )
304        if offspring_positions is not None:
305            offspring_positions = offspring_positions[order]
306
307        # Make eggs
308        eggs = Population.make_eggs(
309            offspring_genomes=offspring_genomes,
310            step=variables.steps,
311            offspring_sexes=offspring_sexes,
312            parental_generations=np.zeros(len(offspring_sexes)),  # TODO replace with working calculation
313            offspring_ancestry=offspring_ancestry,
314            offspring_lineage_id=offspring_lineage_id,
315            offspring_parent_lineage_id=offspring_parent_lineage_id,
316            offspring_positions=offspring_positions,
317        )
318        if self.eggs is None:
319            self.eggs = eggs
320        else:
321            self.eggs += eggs
322
323        if parametermanager.parameters.CARRYING_CAPACITY_EGGS is not None and len(self.eggs) > parametermanager.parameters.CARRYING_CAPACITY_EGGS:
324            indices = np.arange(len(self.eggs))[-parametermanager.parameters.CARRYING_CAPACITY_EGGS :]
325            # TODO biased
326            self.eggs *= indices
327            # Truncated eggs lose their claim on the lattice cells too.
328            if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
329                submodels.lattice.resync_occupancy_from_positions(
330                    self.population.positions, self.eggs.positions
331                )

Generate offspring of reproducing individuals. Initial is set to 0.

def growth(self):
333    def growth(self):
334        # TODO use already scavenged resources to determine growth
335        # max_growth_potential = self.population.phenotypes.extract(ages=self.population.ages, trait_name="grow")
336        # gathered_resources = submodels.resources.scavenge(max_growth_potential)
337        # self.population.sizes += gathered_resources
338        self.population.sizes += 1
def age(self):
340    def age(self):
341        """Increase age of all by one and kill those that surpass age limit.
342        Age denotes the number of full cycles that an individual survived and reproduced.
343        AGE_LIMIT is the maximum number of full cycles an individual can go through.
344        """
345        self.population.ages += 1
346        mask_kill = self.population.ages >= parametermanager.parameters.AGE_LIMIT
347        self._kill(mask_kill=mask_kill, causeofdeath="age_limit")

Increase age of all by one and kill those that surpass age limit. Age denotes the number of full cycles that an individual survived and reproduced. AGE_LIMIT is the maximum number of full cycles an individual can go through.

def hatch(self):
349    def hatch(self):
350        """Turn eggs into living individuals"""
351
352        # If nothing to hatch
353        if self.eggs is None or len(self.eggs) == 0:
354            return
355
356        # If REPRODUCTION_REGULATION is True, only reproduce until MAX_POPULATION_SIZE
357        if parametermanager.parameters.REPRODUCTION_REGULATION:
358            current_population_size = len(self.population)
359            remaining_capacity = resources.capacity - current_population_size
360            # If no remaining capacity, do not reproduce
361            if remaining_capacity < 1:
362                self.eggs = None
363                return
364            elif remaining_capacity < len(self.eggs):
365                indices = variables.rng.choice(len(self.eggs), size=int(remaining_capacity), replace=False)
366                self.eggs *= indices
367
368        # If something to hatch
369        if (
370            (
371                parametermanager.parameters.INCUBATION_PERIOD == -1 and len(self.population) == 0
372            )  # hatch when everyone dead
373            or (parametermanager.parameters.INCUBATION_PERIOD == 0)  # hatch immediately
374            or (
375                parametermanager.parameters.INCUBATION_PERIOD > 0
376                and variables.steps % parametermanager.parameters.INCUBATION_PERIOD == 0
377            )  # hatch with delay
378        ):
379
380            # Lattice mode requires eggs to carry positions, assigned at reproduction
381            # time. If they don't, fail loudly rather than silently corrupt the
382            # lattice with (-1, -1) sentinel positions. (Offspring placement on
383            # the lattice is the next commit.)
384            if parametermanager.parameters.LATTICE_MODE and self.eggs.positions is None:
385                raise NotImplementedError(
386                    "LATTICE_MODE=True requires offspring placement on the lattice, "
387                    "which is not yet wired into reproduction. Use LATTICE_MODE=False "
388                    "for now, or wait for the lattice-reproduction commit."
389                )
390
391            self.eggs.phenotypes = submodels.architect.__call__(self.eggs.genomes)
392            self.population += self.eggs
393            self.eggs = None
394
395            # Sync lattice occupancy after population growth so migration sees
396            # all hatched individuals. No-op when LATTICE_MODE is False.
397            if parametermanager.parameters.LATTICE_MODE and self.population.positions is not None:
398                submodels.lattice.resync_occupancy_from_positions(self.population.positions)

Turn eggs into living individuals