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, resuming=False):
15        self.odir = odir / "gui"
16        self.init_odir()
17        if not resuming:
18            self.init_headers()
19
20    def record(self, population):
21        """Record data that is needed by gui."""
22        # TODO rename INTERVAL_RATE into something more representative; potentially, restructure the recording rates
23        if skip("INTERVAL_RATE") or len(population) == 0:
24            return
25
26        self.write_genotypes(population=population)
27        self.write_phenotypes(population=population)
28
29    def init_headers(self):
30        with open(self.odir / "genotypes.csv", "ab") as f:
31            length = submodels.architect.architecture.length
32            ploidy = submodels.genetics.ploider.ploider.y
33            header0 = list(range(length)) * ploidy
34            header1 = np.repeat(np.arange(ploidy), length)
35            np.savetxt(f, [header0], delimiter=",", fmt="%i")
36            np.savetxt(f, [header1], delimiter=",", fmt="%i")
37
38        with open(self.odir / "phenotypes.csv", "ab") as f:
39            age_limit = submodels.architect.architecture.AGE_LIMIT
40            trait_list = [trait.name for trait in parameterization.traits.values() if trait.evolvable]
41            n_traits = len(trait_list)
42            header0 = np.repeat(trait_list, age_limit)
43            header1 = list(np.arange(age_limit)) * n_traits
44            np.savetxt(f, [header0], delimiter=",", fmt="%s")
45            np.savetxt(f, [header1], delimiter=",", fmt="%i")
46
47    def write_genotypes(self, population: Population):
48        """
49        genotypes.csv | Record allele frequency
50
51        # OUTPUT SPECIFICATION
52        path: /gui/genotypes.csv
53        filetype: csv
54        category: genotype
55        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).
56        trait granularity: population mean
57        time granularity: snapshot
58        frequency parameter: INTERVAL_RATE
59        structure: A float matrix; rows: recordings, columns: genome positions, values: average bit states
60        header: two-row header; first row: genome position, second row: which chromosomal set (0 or 1)
61        """
62        with open(self.odir / "genotypes.csv", "ab") as f:
63            array = population.genomes.flatten().mean(0)
64            np.savetxt(f, [array], delimiter=",", fmt="%1.3e")
65
66    def write_phenotypes(self, population: Population):
67        """
68        phenotypes.csv | Record median phenotype
69
70        # OUTPUT SPECIFICATION
71        path: /gui/genotypes.csv
72        filetype: csv
73        category: phenotype
74        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).
75        trait granularity: population median
76        time granularity: snapshot
77        frequency parameter: INTERVAL_RATE
78        structure: A float matrix; rows: recordings, columns: individual phenotypic traits, values: median trait values
79        header: two-row header; first row: trait name, second row: age
80        """
81        with open(self.odir / "phenotypes.csv", "ab") as f:
82            array = np.median(population.phenotypes.get(), 0)
83            np.savetxt(f, [array], delimiter=",", fmt="%1.3e")