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.parameterization import parametermanager
 12from aegis_sim.utilities.funcs import steps_to_end, skip
 13
 14
 15class FeatherRecorder(Recorder):
 16    def __init__(self, odir: pathlib.Path):
 17        self.odir_genotypes = odir / "snapshots" / "genotypes"
 18        self.odir_phenotypes = odir / "snapshots" / "phenotypes"
 19        self.odir_demography = odir / "snapshots" / "demography"
 20        self.init_dir(self.odir_genotypes)
 21        self.init_dir(self.odir_phenotypes)
 22        self.init_dir(self.odir_demography)
 23
 24    def write(self, population: Population):
 25        """Record demographic, genetic and phenotypic data from the current population."""
 26
 27        # If not final snapshots to be taken, and about to skip or the population is extinct, do not write.
 28        final_snapshots = parametermanager.parameters.SNAPSHOT_FINAL_COUNT > steps_to_end()
 29        if not final_snapshots and (skip("SNAPSHOT_RATE") or len(population) == 0):
 30            return
 31
 32        step = variables.steps
 33
 34        logging.debug(f"Snapshots recorded at step {step}.")
 35
 36        self.write_genotypes(step=step, population=population)
 37        self.write_phenotypes(step=step, population=population)
 38        self.write_demography(step=step, population=population)
 39
 40    def write_genotypes(self, step: int, population: Population):
 41        """
 42
 43        # OUTPUT SPECIFICATION
 44        path: /snapshots/genotypes/{step}.feather
 45        filetype: feather
 46        category: genotype
 47        description: A snapshot of complete binary genomes of all individuals at a certain simulation step.
 48        trait granularity: individual
 49        time granularity: snapshot
 50        frequency parameter: SNAPSHOT_RATE
 51        structure: A bool matrix; rows: individuals, columns: genome positions, values: bit states
 52        header: genome positions
 53        """
 54        if len(population) == 0:
 55            # Empty population can't be reshaped by genomes.flatten(); write empty feather
 56            df_gen = pd.DataFrame()
 57        else:
 58            df_gen = pd.DataFrame(np.array(population.genomes.flatten()))
 59            df_gen.columns = [str(c) for c in df_gen.columns]
 60        df_gen.reset_index(drop=True, inplace=True)
 61        df_gen.to_feather(self.odir_genotypes / f"{step}.feather")
 62
 63    def write_phenotypes(self, step: int, population: Population):
 64        # TODO add more info to columns and rows
 65        """
 66
 67        # OUTPUT SPECIFICATION
 68        path: /snapshots/phenotypes/{step}.feather
 69        filetype: feather
 70        category: phenotype
 71        description: A snapshot of complete intrinsic phenotypes of all individuals at a certain simulation step.
 72        trait granularity: individual
 73        time granularity: snapshot
 74        frequency parameter: SNAPSHOT_RATE
 75        structure: A float matrix; rows: individuals, columns: individual phenotypic traits (depending on which traits are evolvable and what is max lifespan), values: trait values
 76        """
 77        # TODO bugged, wrong header
 78        df_phe = pd.DataFrame(population.phenotypes.get())
 79        df_phe.reset_index(drop=True, inplace=True)
 80        df_phe.columns = [str(c) for c in df_phe.columns]
 81        df_phe.to_feather(self.odir_phenotypes / f"{step}.feather")
 82
 83    def write_demography(self, step: int, population: Population):
 84        """
 85
 86        # OUTPUT SPECIFICATION
 87        path: /snapshots/demography/{step}.feather
 88        filetype: feather
 89        category: demography
 90        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.
 91        trait granularity: individual
 92        time granularity: snapshot
 93        frequency parameter: SNAPSHOT_RATE
 94        structure: A matrix of ints and floats
 95        header: ['ages', 'births', 'birthdays', 'sizes', 'sexes']
 96        """
 97        dem_attrs = [
 98            "ages",
 99            "births",
100            "birthdays",
101            # "generations",
102            "sizes",
103            "sexes",
104        ]
105        demo = {attr: getattr(population, attr) for attr in dem_attrs}
106        df_dem = pd.DataFrame(demo, columns=dem_attrs)
107        df_dem.reset_index(drop=True, inplace=True)
108        df_dem.to_feather(self.odir_demography / f"{step}.feather")