aegis_sim.recording.recordingmanager

  1"""Data recorder
  2
  3Records data generated by the simulation.
  4
  5When thinking about recording additional data, consider that there are three recording methods:
  6    I. Snapshots (record data from the population at a specific step)
  7    II. Flushes (collect data over time then flush)
  8    III. One-time records
  9    IV. Other: TE records
 10"""
 11
 12import pathlib
 13import shutil
 14import logging
 15
 16from aegis_sim import variables
 17
 18from .terecorder import TERecorder
 19from .picklerecorder import PickleRecorder
 20from .popgenstatsrecorder import PopgenStatsRecorder
 21from .intervalrecorder import IntervalRecorder
 22from .flushrecorder import FlushRecorder
 23from .featherrecorder import FeatherRecorder
 24from .phenomaprecorder import PhenomapRecorder
 25from .summaryrecorder import SummaryRecorder
 26from .progressrecorder import ProgressRecorder
 27from .simpleprogressrecorder import SimpleProgressRecorder
 28from .popsizerecorder import PopsizeRecorder
 29from .resourcerecorder import ResourcesRecorder
 30from .ticker import Ticker
 31from .configrecorder import ConfigRecorder
 32from .envdriftmaprecorder import Envdriftmaprecorder
 33from .checkpointrecorder import CheckpointRecorder
 34from .ancestryrecorder import AncestryRecorder
 35from .fastarecorder import FastaRecorder
 36from .vcfrecorder import VCFRecorder
 37from .gvcfrecorder import GVCFRecorder
 38from .lineagerecorder import LineageRecorder
 39from .selectionrecorder import SelectionRecorder
 40from .latticerecorder import LatticeRecorder
 41
 42# TODO write tests
 43
 44
 45class RecordingManager:
 46    """
 47    Container class for various recorders.
 48    Each recorder records a certain type of data.
 49    Most recorders record data as tables, except SummaryRecorder and PickleRecorder which record JSON files and pickles (a binary python format).
 50    Headers and indexes of all tabular files are explicitly recorded.
 51
 52    -----
 53    GUI
 54    AEGIS records a lot of different data.
 55    In brief, AEGIS records
 56    genomic data (population-level allele frequencies and individual-level binary sequences) and
 57    phenotypic data (observed population-level phenotypes and intrinsic individual-level phenotypes),
 58    as well as
 59    derived demographic data (life, death and birth tables),
 60    population genetic data (e.g. effective population size, theta), and
 61    survival analysis data (TE / time-event tables).
 62    Furthermore, it records metadata (e.g. simulation log, processed configuration files) and python pickle files.
 63
 64    Recorded data is distributed in multiple files.
 65    Almost all data are tabular, so each file is a table to which rows are appended as the simulation is running.
 66    The recording rates are frequencies at which rows are added; they are expressed in simulation steps.
 67    """
 68
 69    def init(self, custom_config_path, overwrite):
 70        self.odir = self.make_odir(custom_config_path=custom_config_path, overwrite=overwrite)
 71        self.resuming = False
 72
 73    def init_for_resume(self, custom_config_path):
 74        """Initialize for resume mode — reuse existing output directory, no overwrite."""
 75        output_path = custom_config_path.parent / custom_config_path.stem
 76        if not output_path.exists():
 77            raise FileNotFoundError(
 78                f"Cannot resume: output directory {output_path} does not exist."
 79            )
 80        self.odir = output_path
 81        self.resuming = True
 82
 83    def truncate_for_resume(self, checkpoint_step):
 84        """Truncate output files back to the checkpoint step to avoid duplicate data.
 85
 86        Called after initialize_recorders so self.odir is set. The checkpoint is
 87        saved at the end of run_step for a given step, and all recorders also
 88        write during that same run_step. So the output files already contain data
 89        for checkpoint_step. On resume the sim loop re-executes from checkpoint_step,
 90        so we need to remove data from checkpoint_step onward.
 91
 92        For per-step files (1 line per step, no header), the number of lines to
 93        keep is checkpoint_step - 1 (steps 1 through checkpoint_step-1).
 94
 95        For rate-based files (1 header line + 1 data line every RATE steps),
 96        the number of data lines to keep is (checkpoint_step - 1) // RATE,
 97        plus the header line(s).
 98        """
 99        from aegis_sim.parameterization import parametermanager
100
101        step = checkpoint_step
102
103        # Per-step files: 1 line per step, no header. Keep step-1 lines.
104        per_step_files = [
105            "popsize_before_reproduction.csv",
106            "popsize_after_reproduction.csv",
107            "eggnum_after_reproduction.csv",
108            "resources_before_scavenging.csv",
109            "resources_after_scavenging.csv",
110        ]
111        for fname in per_step_files:
112            self._truncate_file(self.odir / fname, keep_lines=step - 1)
113
114        # Rate-based files with 1 header line.
115        # Data is written when step % RATE == 0 or step == 1.
116        # Number of data lines at step S: 1 (for step 1) + count of multiples of RATE in [2, S-1]
117        # Simpler: lines where skip() returns False for steps 1..S-1
118
119        rate_specs_1header = [
120            ("progress.log", "LOGGING_RATE"),
121        ]
122        for fname, rate_name in rate_specs_1header:
123            rate = getattr(parametermanager.parameters, rate_name)
124            n_data = self._count_recordings(step - 1, rate)
125            self._truncate_file(self.odir / fname, keep_lines=1 + n_data)
126
127        # FlushRecorder spectra: 1 header + data at INTERVAL_RATE
128        spectra_dir = self.odir / "gui" / "spectra"
129        if spectra_dir.exists():
130            rate = parametermanager.parameters.INTERVAL_RATE
131            n_data = self._count_recordings(step - 1, rate)
132            for csv_file in spectra_dir.glob("*.csv"):
133                self._truncate_file(csv_file, keep_lines=1 + n_data)
134
135        # IntervalRecorder: 2 header lines + data at INTERVAL_RATE
136        rate = parametermanager.parameters.INTERVAL_RATE
137        n_data = self._count_recordings(step - 1, rate)
138        for fname in ["gui/genotypes.csv", "gui/phenotypes.csv"]:
139            self._truncate_file(self.odir / fname, keep_lines=2 + n_data)
140
141        # PopgenStatsRecorder: no header, data at POPGENSTATS_RATE
142        popgen_dir = self.odir / "popgen"
143        if popgen_dir.exists():
144            rate = parametermanager.parameters.POPGENSTATS_RATE
145            n_data = self._count_recordings(step - 1, rate)
146            for csv_file in popgen_dir.glob("*.csv"):
147                self._truncate_file(csv_file, keep_lines=n_data)
148
149        # Envdriftmap: no header, writes when step % ENVDRIFT_RATE == 0 (NOT at step 1 unless divisible).
150        # This differs from skip() logic, so we count multiples directly.
151        rate = parametermanager.parameters.ENVDRIFT_RATE
152        if rate > 0:
153            n_data = (step - 1) // rate  # multiples of rate in [1, step-1]
154            self._truncate_file(self.odir / "envdriftmap.csv", keep_lines=n_data)
155
156        # TE files: numbered CSV files in te/ directory.
157        # A new TE file starts at step 1 and every TE_RATE steps.
158        # TE_number increments when a file is flushed (at TE_DURATION offset or final step).
159        # We need to figure out which TE file was in-progress at the checkpoint step
160        # and delete any files started after it.
161        te_rate = parametermanager.parameters.TE_RATE
162        te_duration = parametermanager.parameters.TE_DURATION
163        if te_rate > 0:
164            te_dir = self.odir / "te"
165            if te_dir.exists():
166                self._truncate_te_files(te_dir, step, te_rate, te_duration)
167
168        logging.info(f"Output files truncated to checkpoint step {step}.")
169
170    @staticmethod
171    def _truncate_te_files(te_dir, checkpoint_step, te_rate, te_duration):
172        """Handle TE file truncation on resume.
173
174        TE files are numbered 0.csv, 1.csv, etc. A new collection window opens
175        at step 1 and every TE_RATE steps. The file number increments when the
176        window is flushed. We need to:
177        1. Determine how many complete TE windows finished before checkpoint_step
178        2. Delete any TE files beyond that
179        3. Truncate the in-progress TE file (remove data recorded at/after checkpoint_step)
180        """
181        # Count how many TE windows were fully completed before checkpoint_step.
182        # A window starts at step S where S % te_rate == 0 (or step 1).
183        # It flushes at S + te_duration (or at STEPS_PER_SIMULATION).
184        # Window 0 starts at step 1, flushes at step te_duration (if te_duration < te_rate).
185        # Window i starts at step i*te_rate, flushes at step i*te_rate + te_duration.
186        # A window is "complete" if its flush step < checkpoint_step.
187
188        # Number of complete windows: windows that started AND flushed before checkpoint_step
189        # Window starts: step 1, te_rate, 2*te_rate, 3*te_rate, ...
190        # The window starting at step W flushes at step W + te_duration.
191        # Complete if W + te_duration < checkpoint_step.
192
193        # For simplicity, just count existing TE files and remove those whose
194        # start step >= checkpoint_step.
195        # Window 0 starts at step 1.
196        # Window k starts at step k * te_rate (for k >= 1), or step 1 for k=0.
197
198        existing_files = sorted(te_dir.glob("*.csv"), key=lambda p: int(p.stem))
199        for f in existing_files:
200            file_num = int(f.stem)
201            # Window file_num starts at: step 1 if file_num==0, else file_num * te_rate
202            window_start = 1 if file_num == 0 else file_num * te_rate
203            if window_start >= checkpoint_step:
204                # This window started at or after checkpoint — delete it
205                f.unlink()
206            # If the window started before checkpoint but may contain data from
207            # steps >= checkpoint_step, we need to truncate those lines.
208            # TE files have: 1 header line ("T,E"), then data lines for each
209            # death event recorded during the window. We can't easily map lines
210            # to steps, so we leave partial windows as-is. The resumed sim will
211            # re-open a new window at the appropriate step, and the old partial
212            # data from the interrupted window is acceptable (it's death events
213            # that actually happened).
214
215    @staticmethod
216    def _count_recordings(up_to_step, rate):
217        """Count how many times a recorder with given rate would have written for steps 1..up_to_step.
218
219        Mirrors the skip() logic: always write at step 1, then write at every step divisible by rate.
220        """
221        if rate <= 0 or up_to_step < 1:
222            return 0
223        # Step 1 always records
224        count = 1
225        if up_to_step >= 2:
226            # Multiples of rate in [1, up_to_step] = up_to_step // rate
227            # But step 1 is already counted, so subtract 1 if rate divides 1 (only when rate == 1)
228            count += up_to_step // rate
229            if rate == 1:
230                count -= 1
231        return count
232
233    @staticmethod
234    def _truncate_file(path, keep_lines):
235        """Truncate a file to keep only the first `keep_lines` lines."""
236        if not path.exists():
237            return
238        with open(path, "rb") as f:
239            lines = f.readlines()
240        if len(lines) <= keep_lines:
241            return  # Nothing to truncate
242        with open(path, "wb") as f:
243            f.writelines(lines[:keep_lines])
244
245    def initialize_recorders(self, TICKER_RATE, resuming=False):
246        self.terecorder = TERecorder(odir=self.odir, resuming=resuming)
247        self.picklerecorder = PickleRecorder(odir=self.odir)
248        self.popgenstatsrecorder = PopgenStatsRecorder(odir=self.odir)
249        self.guirecorder = IntervalRecorder(odir=self.odir, resuming=resuming)
250        self.flushrecorder = FlushRecorder(odir=self.odir, resuming=resuming)
251        self.featherrecorder = FeatherRecorder(odir=self.odir)
252        self.phenomaprecorder = PhenomapRecorder(odir=self.odir)
253        self.summaryrecorder = SummaryRecorder(odir=self.odir)
254        self.progressrecorder = ProgressRecorder(odir=self.odir, resuming=resuming)
255        self.simpleprogressrecorder = SimpleProgressRecorder(odir=self.odir)
256        self.ticker = Ticker(odir=self.odir, TICKER_RATE=TICKER_RATE)
257        self.popsizerecorder = PopsizeRecorder(odir=self.odir)
258        self.resourcerecorder = ResourcesRecorder(odir=self.odir)
259        self.configrecorder = ConfigRecorder(odir=self.odir)
260        self.envdriftmaprecorder = Envdriftmaprecorder(odir=self.odir)
261        self.checkpointrecorder = CheckpointRecorder(odir=self.odir)
262        self.ancestryrecorder = AncestryRecorder(odir=self.odir)
263        self.fastarecorder = FastaRecorder(odir=self.odir)
264        self.vcfrecorder = VCFRecorder(odir=self.odir)
265        self.gvcfrecorder = GVCFRecorder(odir=self.odir)
266        self.lineagerecorder = LineageRecorder(odir=self.odir)
267        self.selectionrecorder = SelectionRecorder(odir=self.odir)
268        self.latticerecorder = LatticeRecorder(odir=self.odir)
269
270    #############
271    # UTILITIES #
272    #############
273
274    @staticmethod
275    def make_odir(custom_config_path, overwrite) -> pathlib.Path:
276        output_path = custom_config_path.parent / custom_config_path.stem  # remove .yml
277        is_occupied = output_path.exists() and output_path.is_dir()
278        if is_occupied:
279            if overwrite:
280                shutil.rmtree(output_path)
281            else:
282                raise Exception(f"{output_path} already exists. To overwrite, add flag --overwrite or -o.")
283        return output_path
284
285    @staticmethod
286    def make_subfolders(paths):
287        # TODO relink paths that are now in classes
288        for path in paths:
289            path.mkdir(exist_ok=True, parents=True)
290
291    def is_extinct(self) -> bool:
292        if self.summaryrecorder.extinct:
293            logging.info(f"Population went extinct (at step {variables.steps}).")
294            return True
295        return False