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            header0, header1 = [], []
40            for trait in parameterization.traits.values():
41                if not trait.evolvable:
42                    continue
43                if trait.agespecific:
44                    header0.extend([trait.name] * trait.length)
45                    header1.extend(range(trait.length))
46                else:
47                    header0.append(trait.name)
48                    header1.append(0)
49            np.savetxt(f, [header0], delimiter=",", fmt="%s")
50            np.savetxt(f, [header1], delimiter=",", fmt="%i")
51
52    def write_genotypes(self, population: Population):
53        """
54        genotypes.csv | Record allele frequency
55
56        # OUTPUT SPECIFICATION
57        path: /gui/genotypes.csv
58        filetype: csv
59        category: genotype
60        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).
61        trait granularity: population mean
62        time granularity: snapshot
63        frequency parameter: INTERVAL_RATE
64        structure: A float matrix; rows: recordings, columns: genome positions, values: average bit states
65        header: two-row header; first row: genome position, second row: which chromosomal set (0 or 1)
66        """
67        with open(self.odir / "genotypes.csv", "ab") as f:
68            array = population.genomes.flatten().mean(0)
69            np.savetxt(f, [array], delimiter=",", fmt="%1.3e")
70
71    def write_phenotypes(self, population: Population):
72        """
73        phenotypes.csv | Record median phenotype
74
75        # OUTPUT SPECIFICATION
76        path: /gui/genotypes.csv
77        filetype: csv
78        category: phenotype
79        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).
80        trait granularity: population median
81        time granularity: snapshot
82        frequency parameter: INTERVAL_RATE
83        structure: A float matrix; rows: recordings, columns: individual phenotypic traits, values: median trait values
84        header: two-row header; first row: trait name, second row: age
85        """
86        with open(self.odir / "phenotypes.csv", "ab") as f:
87            array = np.median(population.phenotypes.get(), 0)
88            np.savetxt(f, [array], delimiter=",", fmt="%1.3e")