1import numpy as np
 2import pathlib
 3
 4from aegis_sim.dataclasses.population import Population
 5from aegis_sim import parameterization
 6
 7from aegis_sim.utilities.funcs import skip
 8from .recorder import Recorder
 9from aegis_sim import submodels
10
11
12class IntervalRecorder(Recorder):
13
14    def __init__(self, odir: pathlib.Path):
15        self.odir = odir / "gui"
16        self.init_odir()
17        self.init_headers()
18
19    def record(self, population):
20        """Record data that is needed by gui."""
21        # TODO rename INTERVAL_RATE into something more representative; potentially, restructure the recording rates
22        if skip("INTERVAL_RATE") or len(population) == 0:
23            return
24
25        self.write_genotypes(population=population)
26        self.write_phenotypes(population=population)
27
28    def init_headers(self):
29        with open(self.odir / "genotypes.csv", "ab") as f:
30            length = submodels.architect.architecture.length
31            ploidy = submodels.genetics.ploider.ploider.y
32            header0 = list(range(length)) * ploidy
33            header1 = np.repeat(np.arange(ploidy), length)
34            np.savetxt(f, [header0], delimiter=",", fmt="%i")
35            np.savetxt(f, [header1], delimiter=",", fmt="%i")
36
37        with open(self.odir / "phenotypes.csv", "ab") as f:
38            age_limit = submodels.architect.architecture.AGE_LIMIT
39            trait_list = [trait.name for trait in parameterization.traits.values() if trait.evolvable]
40            n_traits = len(trait_list)
41            header0 = np.repeat(trait_list, age_limit)
42            header1 = list(np.arange(age_limit)) * n_traits
43            np.savetxt(f, [header0], delimiter=",", fmt="%s")
44            np.savetxt(f, [header1], delimiter=",", fmt="%i")
45
46    def write_genotypes(self, population: Population):
47        """
48        genotypes.csv | Record allele frequency
49
50        # OUTPUT SPECIFICATION
51        path: /gui/genotypes.csv
52        filetype: csv
53        category: genotype
54        description: A table of allele frequencies (frequency of 1's in the population for each site) across simulation intervals. Columns are sites, rows are simulation intervals (spanning INTERVAL_RATE steps).
55        trait granularity: population mean
56        time granularity: snapshot
57        frequency parameter: INTERVAL_RATE
58        structure: A float matrix; rows: recordings, columns: genome positions, values: average bit states
59        header: two-row header; first row: genome position, second row: which chromosomal set (0 or 1)
60        """
61        with open(self.odir / "genotypes.csv", "ab") as f:
62            array = population.genomes.flatten().mean(0)
63            np.savetxt(f, [array], delimiter=",", fmt="%1.3e")
64
65    def write_phenotypes(self, population: Population):
66        """
67        phenotypes.csv | Record median phenotype
68
69        # OUTPUT SPECIFICATION
70        path: /gui/genotypes.csv
71        filetype: csv
72        category: phenotype
73        description: A table of median intrinsic phenotypes (median phenotype rate for each trait at each age) across simulation intervals. Columns are traits, rows are simulation intervals (spanning INTERVAL_RATE steps).
74        trait granularity: population median
75        time granularity: snapshot
76        frequency parameter: INTERVAL_RATE
77        structure: A float matrix; rows: recordings, columns: individual phenotypic traits, values: median trait values
78        header: two-row header; first row: trait name, second row: age
79        """
80        with open(self.odir / "phenotypes.csv", "ab") as f:
81            array = np.median(population.phenotypes.get(), 0)
82            np.savetxt(f, [array], delimiter=",", fmt="%1.3e")