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):
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        self.init_headers()
34
35    def collect(self, key, ages):
36        """Add data into memory which will be recorded later."""
37        self.collection[key] += np.bincount(ages, minlength=self.n_ages)
38
39    def flush(self):
40        """Record data that has been collected over time."""
41        # spectra/*.csv | Age distribution of various subpopulations (e.g. population that died of genetic causes)
42
43        if skip("INTERVAL_RATE"):
44            return
45
46        for key, val in self.collection.items():
47            self.write_age_at(filename=key, collected_values=val)
48            # with open(self.odir / f"{key}.csv", "ab") as f:
49            #     array = np.array(val)
50            #     np.savetxt(f, [array], delimiter=",", fmt="%i")
51
52        # Reinitialize the collection
53        self.collection = copy.deepcopy(self._collection)
54
55    def write_age_at(self, filename, collected_values):
56        """
57
58        # OUTPUT SPECIFICATION
59        path: /gui/spectra/age_at_{cause}.csv
60        filetype: csv
61        category: demography
62        description: Total number of deaths by age and cause of death, within a simulation interval.
63        trait granularity: population count
64        time granularity: interval
65        frequency parameter: INTERVAL_RATE
66        structure: An int matrix.
67        header: list of ages from 0 to AGE_LIMIT (inclusive)
68        """
69        with open(self.odir / f"{filename}.csv", "ab") as f:
70            array = np.array(collected_values)
71            np.savetxt(f, [array], delimiter=",", fmt="%i")
72
73    def init_headers(self):
74        for key in self._collection.keys():
75            with open(self.odir / f"{key}.csv", "ab") as f:
76                array = np.arange(self.n_ages)
77                np.savetxt(f, [array], delimiter=",", fmt="%i")