1import copy
 2import numpy as np
 3from aegis_sim.constants import VALID_CAUSES_OF_DEATH
 4import pathlib
 5from aegis_sim.utilities.funcs import skip
 6
 7from aegis_sim.parameterization import parametermanager
 8from .recorder import Recorder
 9
10
11class FlushRecorder(Recorder):
12    """
13
14    Records collections.
15    """
16
17    def __init__(self, odir: pathlib.Path, resuming=False):
18
19        self.odir = odir / "gui" / "spectra"
20        self.init_odir()
21
22        self.n_ages = parametermanager.parameters.AGE_LIMIT + 1
23
24        self._collection = {
25            "age_at_birth": [0] * self.n_ages,
26            "additive_age_structure": [0] * self.n_ages,
27        }
28
29        self._collection.update({f"age_at_{causeofdeath}": [0] * self.n_ages for causeofdeath in VALID_CAUSES_OF_DEATH})
30
31        self.collection = copy.deepcopy(self._collection)
32
33        if not resuming:
34            self.init_headers()
35
36    def collect(self, key, ages):
37        """Add data into memory which will be recorded later."""
38        self.collection[key] += np.bincount(ages, minlength=self.n_ages)
39
40    def flush(self):
41        """Record data that has been collected over time."""
42        # spectra/*.csv | Age distribution of various subpopulations (e.g. population that died of genetic causes)
43
44        if skip("INTERVAL_RATE"):
45            return
46
47        for key, val in self.collection.items():
48            self.write_age_at(filename=key, collected_values=val)
49            # with open(self.odir / f"{key}.csv", "ab") as f:
50            #     array = np.array(val)
51            #     np.savetxt(f, [array], delimiter=",", fmt="%i")
52
53        # Reinitialize the collection
54        self.collection = copy.deepcopy(self._collection)
55
56    def write_age_at(self, filename, collected_values):
57        """
58
59        # OUTPUT SPECIFICATION
60        path: /gui/spectra/age_at_{cause}.csv
61        filetype: csv
62        category: demography
63        description: Total number of deaths by age and cause of death, within a simulation interval.
64        trait granularity: population count
65        time granularity: interval
66        frequency parameter: INTERVAL_RATE
67        structure: An int matrix.
68        header: list of ages from 0 to AGE_LIMIT (inclusive)
69        """
70        with open(self.odir / f"{filename}.csv", "ab") as f:
71            array = np.array(collected_values)
72            np.savetxt(f, [array], delimiter=",", fmt="%i")
73
74    def init_headers(self):
75        for key in self._collection.keys():
76            with open(self.odir / f"{key}.csv", "ab") as f:
77                array = np.arange(self.n_ages)
78                np.savetxt(f, [array], delimiter=",", fmt="%i")