1import logging
  2import pandas as pd
  3import numpy as np
  4
  5import pathlib
  6
  7from aegis_sim.dataclasses.population import Population
  8from .recorder import Recorder
  9from aegis_sim import variables
 10
 11from aegis_sim import parameterization
 12from aegis_sim.parameterization import parametermanager
 13from aegis_sim.constants import GENETIC_TRAITS
 14from aegis_sim.utilities.funcs import steps_to_end, skip
 15
 16
 17class FeatherRecorder(Recorder):
 18    def __init__(self, odir: pathlib.Path):
 19        self.odir_genotypes = odir / "snapshots" / "genotypes"
 20        self.odir_phenotypes = odir / "snapshots" / "phenotypes"
 21        self.odir_demography = odir / "snapshots" / "demography"
 22        self.init_dir(self.odir_genotypes)
 23        self.init_dir(self.odir_phenotypes)
 24        self.init_dir(self.odir_demography)
 25
 26    def write(self, population: Population):
 27        """Record demographic, genetic and phenotypic data from the current population."""
 28
 29        # If not final snapshots to be taken, and about to skip or the population is extinct, do not write.
 30        final_snapshots = parametermanager.parameters.SNAPSHOT_FINAL_COUNT > steps_to_end()
 31        if not final_snapshots and (skip("SNAPSHOT_RATE") or len(population) == 0):
 32            return
 33
 34        step = variables.steps
 35
 36        logging.debug(f"Snapshots recorded at step {step}.")
 37
 38        self.write_genotypes(step=step, population=population)
 39        self.write_phenotypes(step=step, population=population)
 40        self.write_demography(step=step, population=population)
 41
 42    def write_genotypes(self, step: int, population: Population):
 43        """
 44
 45        # OUTPUT SPECIFICATION
 46        path: /snapshots/genotypes/{step}.feather
 47        filetype: feather
 48        category: genotype
 49        description: A snapshot of complete binary genomes of all individuals at a certain simulation step.
 50        trait granularity: individual
 51        time granularity: snapshot
 52        frequency parameter: SNAPSHOT_RATE
 53        structure: A bool matrix; rows: individuals, columns: genome positions, values: bit states
 54        header: genome positions
 55        """
 56        if len(population) == 0:
 57            # Empty population can't be reshaped by genomes.flatten(); write empty feather
 58            df_gen = pd.DataFrame()
 59        else:
 60            df_gen = pd.DataFrame(np.array(population.genomes.flatten()))
 61            df_gen.columns = [str(c) for c in df_gen.columns]
 62        df_gen.reset_index(drop=True, inplace=True)
 63        df_gen.to_feather(self.odir_genotypes / f"{step}.feather")
 64
 65    def write_phenotypes(self, step: int, population: Population):
 66        """
 67
 68        # OUTPUT SPECIFICATION
 69        path: /snapshots/phenotypes/{step}.feather
 70        filetype: feather
 71        category: phenotype
 72        description: A snapshot of complete intrinsic phenotypes of all individuals at a certain simulation step.
 73        trait granularity: individual
 74        time granularity: snapshot
 75        frequency parameter: SNAPSHOT_RATE
 76        structure: A float matrix; rows: individuals, columns: {trait}_{age} for each evolvable trait, values: trait values
 77        header: {trait}_{age} e.g. surv_0, surv_1, ..., repr_0, repr_1, ...
 78        """
 79        df_phe = pd.DataFrame(population.phenotypes.get())
 80        df_phe.reset_index(drop=True, inplace=True)
 81        df_phe.columns = self._phenotype_columns()
 82        df_phe.to_feather(self.odir_phenotypes / f"{step}.feather")
 83
 84    @staticmethod
 85    def _phenotype_columns():
 86        AGE_LIMIT = parametermanager.parameters.AGE_LIMIT
 87        cols = []
 88        for trait_name in GENETIC_TRAITS:
 89            trait = parameterization.traits[trait_name]
 90            if not trait.evolvable:
 91                continue
 92            if trait.agespecific:
 93                cols.extend(f"{trait_name}_{age}" for age in range(AGE_LIMIT))
 94            else:
 95                cols.append(trait_name)
 96        return cols
 97
 98    def write_demography(self, step: int, population: Population):
 99        """
100
101        # OUTPUT SPECIFICATION
102        path: /snapshots/demography/{step}.feather
103        filetype: feather
104        category: demography
105        description: A recording of life history metrics (age, number of births given, step at which born, current size, sex) of all individuals until a certain simulation step.
106        trait granularity: individual
107        time granularity: snapshot
108        frequency parameter: SNAPSHOT_RATE
109        structure: A matrix of ints and floats
110        header: ['ages', 'births', 'birthdays', 'sizes', 'sexes']
111        """
112        dem_attrs = [
113            "ages",
114            "births",
115            "birthdays",
116            # "generations",
117            "sizes",
118            "sexes",
119        ]
120        demo = {attr: getattr(population, attr) for attr in dem_attrs}
121        df_dem = pd.DataFrame(demo, columns=dem_attrs)
122        df_dem.reset_index(drop=True, inplace=True)
123        df_dem.to_feather(self.odir_demography / f"{step}.feather")