aegis_sim.dataclasses.population

  1import numpy as np
  2import pickle
  3import pathlib
  4
  5from aegis_sim.dataclasses.genomes import Genomes
  6from aegis_sim.dataclasses.phenotypes import Phenotypes
  7from aegis_sim import submodels
  8
  9
 10class Population:
 11    """Population data
 12
 13    Contains demographic, genetic and phenotypic data of living individuals.
 14    """
 15
 16    attrs = (
 17        "genomes",
 18        "ages",
 19        "births",
 20        "birthdays",
 21        "generations",
 22        "phenotypes",
 23        "infection",
 24        "sizes",
 25        "sexes",
 26    )
 27
 28    def __init__(
 29        self,
 30        genomes: Genomes,
 31        ages,
 32        births,
 33        birthdays,
 34        phenotypes: Phenotypes,
 35        infection,
 36        sizes,
 37        sexes,
 38        generations=None,
 39    ):
 40        self.genomes = genomes
 41        self.ages = ages
 42        self.births = births
 43        self.birthdays = birthdays
 44        self.phenotypes = phenotypes
 45        self.infection = infection
 46        self.sizes = sizes
 47        self.sexes = sexes
 48        self.generations = generations
 49
 50        assert isinstance(phenotypes, Phenotypes)
 51
 52        if not (
 53            len(genomes)
 54            == len(ages)
 55            == len(births)
 56            == len(birthdays)
 57            == len(phenotypes)
 58            == len(infection)
 59            == len(sizes)
 60            == len(sexes)
 61            # == len(generations)
 62        ):
 63            raise ValueError("Population attributes must have equal length")
 64
 65    def __len__(self):
 66        """Return the number of living individuals."""
 67        return len(self.genomes)
 68
 69    def __getitem__(self, index):
 70        """Return a subpopulation."""
 71        return Population(
 72            genomes=self.genomes.get(individuals=index),
 73            ages=self.ages[index],
 74            births=self.births[index],
 75            birthdays=self.birthdays[index],
 76            phenotypes=self.phenotypes.get(individuals=index),
 77            infection=self.infection[index],
 78            sizes=self.sizes[index],
 79            sexes=self.sexes[index],
 80            generations=self.generations[index] if self.generations is not None else None,
 81        )
 82
 83    def __imul__(self, index):
 84        """Redefine itself as its own subpopulation."""
 85        for attr in self.attrs:
 86            if attr == "genomes":
 87                self.genomes.keep(individuals=index)
 88            elif attr == "phenotypes":
 89                self.phenotypes.keep(individuals=index)
 90            elif attr == "generations":
 91                self.generations = None
 92            else:
 93                setattr(self, attr, getattr(self, attr)[index])
 94        return self
 95
 96    def __iadd__(self, population):
 97        """Merge with another population."""
 98
 99        for attr in self.attrs:
100            if attr == "genomes":
101                self.genomes.add(population.genomes)
102            elif attr == "phenotypes":
103                assert isinstance(population.phenotypes, Phenotypes)
104                self.phenotypes.add(population.phenotypes)
105            elif attr == "generations":
106                self.generations = None
107            else:
108                val = np.concatenate([getattr(self, attr), getattr(population, attr)])
109                setattr(self, attr, val)
110        return self
111
112    # def shuffle(self):
113    #     order = np.random.arange(len(self))
114    #     np.random.shuffle(order)
115    #     self *= order
116
117    @staticmethod
118    def load_pickle_from(path: pathlib.Path):
119        assert path.exists(), f"pickle_path {path} does not exist"
120        with open(path, "rb") as file_:
121            return pickle.load(file_)
122
123    def save_pickle_to(self, path):
124        with open(path, "wb") as file_:
125            pickle.dump(self, file_)
126
127    @staticmethod
128    def initialize(n, AGE_LIMIT):
129        genomes = Genomes(submodels.architect.architecture.init_genome_array(n))
130        ages = np.random.randint(low=0, high=AGE_LIMIT, size=n, dtype=np.int32)
131        births = np.zeros(n, dtype=np.int32)
132        birthdays = np.zeros(n, dtype=np.int32)
133        # generations = np.zeros(n, dtype=np.int32)
134        generations = None
135
136        phenotypes = submodels.architect.__call__(genomes)
137        assert isinstance(phenotypes, Phenotypes)
138
139        infection = np.zeros(n, dtype=np.int32)
140        sizes = np.zeros(n, dtype=np.float32)
141        sexes = submodels.sexsystem.get_sex(n)
142        return Population(
143            genomes=genomes,
144            ages=ages,
145            births=births,
146            birthdays=birthdays,
147            generations=generations,
148            phenotypes=phenotypes,
149            infection=infection,
150            sizes=sizes,
151            sexes=sexes,
152        )
153
154    @staticmethod
155    def make_eggs(offspring_genomes: Genomes, step, offspring_sexes, parental_generations):
156        n = len(offspring_genomes)
157        eggs = Population(
158            genomes=offspring_genomes,
159            ages=np.zeros(n, dtype=np.int32),
160            births=np.zeros(n, dtype=np.int32),
161            birthdays=np.zeros(n, dtype=np.int32) + step,
162            # generations=parental_generations + 1,
163            generations=None,
164            # phenotypes=submodels.architect.__call__(offspring_genomes), # Do not compute phenotypes until eggs are laid! Why? Because it is computationally expensive.
165            phenotypes=Phenotypes.init_phenotype_array(n),
166            infection=np.zeros(n, dtype=np.int32),
167            sizes=np.zeros(n, dtype=np.float32),
168            sexes=offspring_sexes,
169        )
170        return eggs