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        df_gen = pd.DataFrame(np.array(population.genomes.flatten()))
 55        df_gen.reset_index(drop=True, inplace=True)
 56        df_gen.columns = [str(c) for c in df_gen.columns]
 57        df_gen.to_feather(self.odir_genotypes / f"{step}.feather")
 58
 59    def write_phenotypes(self, step: int, population: Population):
 60        # TODO add more info to columns and rows
 61        """
 62
 63        # OUTPUT SPECIFICATION
 64        path: /snapshots/phenotypes/{step}.feather
 65        filetype: feather
 66        category: phenotype
 67        description: A snapshot of complete intrinsic phenotypes of all individuals at a certain simulation step.
 68        trait granularity: individual
 69        time granularity: snapshot
 70        frequency parameter: SNAPSHOT_RATE
 71        structure: A float matrix; rows: individuals, columns: individual phenotypic traits (depending on which traits are evolvable and what is max lifespan), values: trait values
 72        """
 73        # TODO bugged, wrong header
 74        df_phe = pd.DataFrame(population.phenotypes.get())
 75        df_phe.reset_index(drop=True, inplace=True)
 76        df_phe.columns = [str(c) for c in df_phe.columns]
 77        df_phe.to_feather(self.odir_phenotypes / f"{step}.feather")
 78
 79    def write_demography(self, step: int, population: Population):
 80        """
 81
 82        # OUTPUT SPECIFICATION
 83        path: /snapshots/demography/{step}.feather
 84        filetype: feather
 85        category: demography
 86        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.
 87        trait granularity: individual
 88        time granularity: snapshot
 89        frequency parameter: SNAPSHOT_RATE
 90        structure: A matrix of ints and floats
 91        header: ['ages', 'births', 'birthdays', 'sizes', 'sexes']
 92        """
 93        dem_attrs = [
 94            "ages",
 95            "births",
 96            "birthdays",
 97            # "generations",
 98            "sizes",
 99            "sexes",
100        ]
101        demo = {attr: getattr(population, attr) for attr in dem_attrs}
102        df_dem = pd.DataFrame(demo, columns=dem_attrs)
103        df_dem.reset_index(drop=True, inplace=True)
104        df_dem.to_feather(self.odir_demography / f"{step}.feather")