aegis_sim.recording.lineagerecorder

  1import logging
  2import pathlib
  3
  4import numpy as np
  5
  6from .recorder import Recorder
  7from aegis_sim import variables
  8from aegis_sim.parameterization import parametermanager
  9
 10
 11HEADER = "step,lineage_id,parent_lineage_id\n"
 12DEATHS_HEADER = "step,lineage_id,cause\n"
 13
 14
 15class LineageRecorder(Recorder):
 16    """Write a growing per-birth log to /lineage/births.csv.
 17
 18    One row per individual ever created: (step at which it was born,
 19    its unique lineage_id, its parent's lineage_id). Initial-population
 20    individuals are written at step 0 with parent_lineage_id=-1.
 21
 22    Triggered inline from bioreactor.reproduction() and from the run()
 23    entry point — NOT from the standard end-of-step recorder cycle, because
 24    it needs to see new births as they happen, not snapshots of the alive
 25    population.
 26
 27    LINEAGE_RATE controls flush cadence to disk (every Nth flush call);
 28    every birth is still recorded regardless of the rate.
 29    """
 30
 31    def __init__(self, odir: pathlib.Path):
 32        self.odir = odir / "lineage"
 33        self.init_odir()
 34        self._file_path = self.odir / "births.csv"
 35        self._deaths_path = self.odir / "deaths.csv"
 36        self._file = None
 37        self._deaths_file = None
 38        self._writes_since_flush = 0
 39
 40    def _ensure_open(self):
 41        if self._file is not None:
 42            return
 43        # Write mode the first time (fresh run); open in append mode otherwise.
 44        # We detect a fresh run by checking whether the header is already present.
 45        mode = "a" if self._file_path.exists() and self._file_path.stat().st_size > 0 else "w"
 46        self._file = open(self._file_path, mode)
 47        if mode == "w":
 48            self._file.write(HEADER)
 49
 50    def _ensure_deaths_open(self):
 51        if self._deaths_file is not None:
 52            return
 53        mode = "a" if self._deaths_path.exists() and self._deaths_path.stat().st_size > 0 else "w"
 54        self._deaths_file = open(self._deaths_path, mode)
 55        if mode == "w":
 56            self._deaths_file.write(DEATHS_HEADER)
 57
 58    def write_initial(self, lineage_ids):
 59        if parametermanager.parameters.LINEAGE_RATE <= 0:
 60            return
 61        if lineage_ids is None or len(lineage_ids) == 0:
 62            return
 63        self._ensure_open()
 64        for lid in lineage_ids:
 65            self._file.write(f"0,{int(lid)},-1\n")
 66        self._maybe_flush()
 67        logging.debug(f"lineage initial: wrote {len(lineage_ids)} rows at step 0.")
 68
 69    def write_births(self, parent_lineage_ids, child_lineage_ids, step):
 70        if parametermanager.parameters.LINEAGE_RATE <= 0:
 71            return
 72        if child_lineage_ids is None or len(child_lineage_ids) == 0:
 73            return
 74        assert len(parent_lineage_ids) == len(child_lineage_ids)
 75        self._ensure_open()
 76        parent_arr = np.asarray(parent_lineage_ids, dtype=np.int64)
 77        child_arr = np.asarray(child_lineage_ids, dtype=np.int64)
 78        for pid, cid in zip(parent_arr.tolist(), child_arr.tolist()):
 79            self._file.write(f"{step},{cid},{pid}\n")
 80        self._maybe_flush()
 81
 82    def write_deaths(self, lineage_ids, causeofdeath, step):
 83        if parametermanager.parameters.LINEAGE_RATE <= 0:
 84            return
 85        if lineage_ids is None or len(lineage_ids) == 0:
 86            return
 87        self._ensure_deaths_open()
 88        for lid in np.asarray(lineage_ids, dtype=np.int64).tolist():
 89            self._deaths_file.write(f"{step},{lid},{causeofdeath}\n")
 90        # Reuse the same flush counter as births.
 91        self._maybe_flush()
 92
 93    def _maybe_flush(self):
 94        self._writes_since_flush += 1
 95        rate = parametermanager.parameters.LINEAGE_RATE
 96        if rate > 0 and self._writes_since_flush >= rate:
 97            if self._file is not None:
 98                self._file.flush()
 99            if self._deaths_file is not None:
100                self._deaths_file.flush()
101            self._writes_since_flush = 0
102
103    def close(self):
104        if self._file is not None:
105            self._file.flush()
106            self._file.close()
107            self._file = None
108        if self._deaths_file is not None:
109            self._deaths_file.flush()
110            self._deaths_file.close()
111            self._deaths_file = None
112
113    # Compatibility no-op so the recorder can be safely added to the standard
114    # end-of-step recording cycle if anyone wires it in there (currently we
115    # don't — see module docstring).
116    def write(self, population):
117        return