aegis_sim

  1import logging
  2import pathlib
  3
  4from aegis_sim.dataclasses.population import Population
  5from aegis_sim.bioreactor import Bioreactor
  6from aegis_sim import variables, submodels, parameterization
  7from aegis_sim.parameterization import parametermanager
  8from aegis_sim.recording import recordingmanager
  9
 10
 11def run(custom_config_path, pickle_path, overwrite, custom_input_params, resume_path=None, extend_steps=None):
 12    if resume_path is not None:
 13        odir = pathlib.Path(resume_path)
 14        checkpoint_file = odir / "checkpoint"
 15
 16        if not odir.exists():
 17            # No output dir yet → fresh run
 18            logging.info(f"No output directory at {odir}, starting fresh run.")
 19            init(custom_config_path, overwrite=False, pickle_path=pickle_path, custom_input_params=custom_input_params)
 20            population = (
 21                Population.initialize(
 22                    n=parametermanager.parameters.INITIAL_POPULATION_SIZE,
 23                    AGE_LIMIT=parametermanager.parameters.AGE_LIMIT,
 24                )
 25                if pickle_path is None
 26                else Population.load_pickle_from(pickle_path)
 27            )
 28            eggs = None
 29        elif not checkpoint_file.exists():
 30            # Output dir exists but no checkpoint → error
 31            raise FileNotFoundError(
 32                f"Output directory {odir} exists but contains no checkpoint file. "
 33                f"Cannot resume. Use -o to overwrite, or delete the directory."
 34            )
 35        else:
 36            # Output dir + checkpoint → resume
 37            population, eggs = init_resume(resume_path, extend_steps=extend_steps)
 38    else:
 39        init(custom_config_path, overwrite, pickle_path, custom_input_params)
 40        population = (
 41            Population.initialize(
 42                n=parametermanager.parameters.INITIAL_POPULATION_SIZE,
 43                AGE_LIMIT=parametermanager.parameters.AGE_LIMIT,
 44            )
 45            if pickle_path is None
 46            else Population.load_pickle_from(pickle_path)
 47        )
 48        eggs = None
 49        population = _seed_introgression(population)
 50
 51        if parametermanager.parameters.LINEAGE_TRACING and population.lineage_id is not None:
 52            recordingmanager.lineagerecorder.write_initial(population.lineage_id)
 53
 54    bioreactor = Bioreactor(population)
 55    bioreactor.eggs = eggs
 56    sim(bioreactor=bioreactor)
 57
 58
 59def init(custom_config_path, overwrite=False, pickle_path=None, custom_input_params={}):
 60    """
 61    When testing aegis, initialize all modules using this function, e.g.
 62
 63    import aegis_sim
 64    aegis_sim.init("_.yml")
 65
 66    And then you can safely import any module.
 67    """
 68
 69    custom_config_path = pathlib.Path(custom_config_path)
 70
 71    parametermanager.init(
 72        custom_config_path=custom_config_path,
 73        custom_input_params=custom_input_params,
 74    )
 75    variables.init(
 76        variables,
 77        custom_config_path=custom_config_path,
 78        pickle_path=pickle_path,
 79        RANDOM_SEED=parametermanager.parameters.RANDOM_SEED,
 80    )
 81    parameterization.init_traits(parameterization)
 82    submodels.init(submodels, parametermanager=parametermanager)
 83
 84    recordingmanager.init(custom_config_path, overwrite)
 85    recordingmanager.initialize_recorders(TICKER_RATE=parametermanager.parameters.TICKER_RATE)
 86
 87    if (
 88        parametermanager.parameters.LINEAGE_TRACING
 89        and parametermanager.parameters.REPRODUCTION_MODE != "asexual"
 90    ):
 91        logging.warning(
 92            "LINEAGE_TRACING is True but REPRODUCTION_MODE is %r; lineage IDs will be "
 93            "assigned to the initial population only and NOT propagated to offspring. "
 94            "Sexual lineage tracing is not yet implemented.",
 95            parametermanager.parameters.REPRODUCTION_MODE,
 96        )
 97
 98
 99def init_resume(resume_path, extend_steps=None):
100    """Initialize all modules from the latest checkpoint in the given output directory.
101
102    Args:
103        resume_path: Path to the output directory containing the checkpoint.
104        extend_steps: If set, override STEPS_PER_SIMULATION to extend the run.
105    """
106    from aegis_sim.checkpoint import Checkpoint
107
108    odir = pathlib.Path(resume_path)
109    checkpoint_path = Checkpoint.find_latest(odir)
110    checkpoint = Checkpoint.load(checkpoint_path)
111
112    # Restore parameters from checkpoint config
113    parametermanager.init_from_config(checkpoint.final_config, checkpoint.custom_config_path)
114
115    # Apply --extend override if provided
116    if extend_steps is not None:
117        if extend_steps <= checkpoint.step:
118            raise ValueError(
119                f"--extend {extend_steps} must be greater than checkpoint step {checkpoint.step}"
120            )
121        parametermanager.parameters.STEPS_PER_SIMULATION = extend_steps
122        parametermanager.final_config["STEPS_PER_SIMULATION"] = extend_steps
123        logging.info(f"Extending simulation to {extend_steps} steps (was {checkpoint.final_config['STEPS_PER_SIMULATION']}).")
124
125    # Restore variables (step, RNG state)
126    variables.restore_from_checkpoint(variables, checkpoint)
127
128    # Re-init traits and submodels
129    parameterization.init_traits(parameterization)
130    submodels.init(submodels, parametermanager=parametermanager)
131
132    # Restore envdrift map if it was active
133    if checkpoint.envdrift_map is not None:
134        submodels.architect.envdrift.map = checkpoint.envdrift_map
135
136    # Restore predator population size
137    submodels.predation.N = checkpoint.predator_population_size
138
139    # Restore resource capacity
140    from aegis_sim.submodels.resources.resources import resources
141    resources.capacity = checkpoint.resource_capacity
142
143    # Init recording in append mode (don't overwrite, don't write headers)
144    recordingmanager.init_for_resume(checkpoint.custom_config_path)
145    recordingmanager.initialize_recorders(
146        TICKER_RATE=parametermanager.parameters.TICKER_RATE,
147        resuming=True,
148    )
149
150    # Truncate output files to remove data recorded after the checkpoint step
151    recordingmanager.truncate_for_resume(checkpoint.step)
152
153    # Update TE recorder's file counter after truncation may have deleted files
154    te_dir = recordingmanager.odir / "te"
155    if te_dir.exists():
156        remaining = list(te_dir.glob("*.csv"))
157        recordingmanager.terecorder.TE_number = len(remaining)
158
159    return checkpoint.population, checkpoint.eggs
160
161
162def _seed_introgression(population: Population) -> Population:
163    """If INTROGRESSION_SOURCE and INTROGRESSION_SEEDS are set, load pop A from pickle,
164    sample the requested number of individuals, mark their ancestry as True (introgressed),
165    set pop B ancestry to False (native), and merge into one population.
166    Returns population unchanged if introgression is not configured.
167    """
168    import numpy as np
169
170    n_seeds = parametermanager.parameters.INTROGRESSION_SEEDS
171    source_path = parametermanager.parameters.INTROGRESSION_SOURCE
172
173    if n_seeds == 0 or source_path is None:
174        return population
175
176    source_path = pathlib.Path(source_path)
177    pop_a = Population.load_pickle_from(source_path)
178
179    if n_seeds > len(pop_a):
180        logging.warning(
181            f"INTROGRESSION_SEEDS={n_seeds} exceeds source population size {len(pop_a)}; "
182            f"clamping to {len(pop_a)}"
183        )
184        n_seeds = len(pop_a)
185
186    # Sample n_seeds individuals from pop A
187    indices = variables.rng.choice(len(pop_a), size=n_seeds, replace=False)
188    pop_a *= indices
189
190    genome_shape = pop_a.genomes.array.shape  # (n_seeds, ploidy, n_loci, bpl)
191
192    # Mark all pop A seeds as introgressed
193    pop_a.ancestry = np.ones(genome_shape, dtype=np.bool_)
194
195    # Mark all pop B individuals as native
196    pop_b_shape = population.genomes.array.shape
197    population.ancestry = np.zeros(pop_b_shape, dtype=np.bool_)
198
199    logging.info(
200        f"Introgression: seeding {n_seeds} individuals from {source_path} into population of {len(population)}"
201    )
202
203    population += pop_a
204    return population
205
206
207def sim(bioreactor):
208    # presim
209    recordingmanager.configrecorder.write_final_config_file(parametermanager.final_config)
210    recordingmanager.ticker.start_process()
211    ticker_pid = recordingmanager.ticker.pid
212    assert ticker_pid is not None
213    recordingmanager.summaryrecorder.write_input_summary(ticker_pid=recordingmanager.ticker.pid)
214    # TODO hacky solution of decrementing and incrementing steps
215    variables.steps -= 1
216    recordingmanager.featherrecorder.write(bioreactor.population)
217    variables.steps += 1
218
219    # sim
220    recordingmanager.phenomaprecorder.write()
221
222    # Write initial checkpoint before the loop so there's always something to resume from
223    if parametermanager.parameters.CHECKPOINT_RATE > 0:
224        from aegis_sim.checkpoint import Checkpoint
225        initial_cp = Checkpoint.capture(bioreactor.population, bioreactor.eggs, variables, submodels, parametermanager)
226        initial_cp.save(recordingmanager.checkpointrecorder.checkpoint_path)
227        logging.debug("Initial checkpoint saved before sim loop.")
228
229    while (variables.steps <= parametermanager.parameters.STEPS_PER_SIMULATION) and not recordingmanager.is_extinct():
230        recordingmanager.progressrecorder.write(len(bioreactor.population))
231        recordingmanager.simpleprogressrecorder.write()
232        bioreactor.run_step()
233        variables.steps += 1
234
235    # postsim
236    recordingmanager.popsizerecorder.flush_all()
237    recordingmanager.resourcerecorder.flush_all()
238    recordingmanager.summaryrecorder.write_output_summary()
239    logging.info("Simulation finished.")
240    recordingmanager.ticker.stop_process()