1"""Lattice spatial snapshot recorder.
 2
 3Writes a per-individual CSV snapshot of the lattice state at the configured
 4rate. Columns: step, q, r, age, sex, lineage_id, ancestry_fraction.
 5
 6`runs/lattice_animate.py` consumes the snapshot files in a sim's
 7`/lattice/` directory and produces a PNG montage (selected steps) or an
 8animated GIF.
 9
10No-op when LATTICE_MODE is False or LATTICE_RECORD_RATE <= 0.
11"""
12
13import logging
14import pathlib
15
16import numpy as np
17
18from .recorder import Recorder
19from aegis_sim import variables
20from aegis_sim.parameterization import parametermanager
21from aegis_sim.utilities.funcs import skip
22
23
24HEADER = "step,q,r,age,sex,lineage_id,ancestry_fraction\n"
25
26
27class LatticeRecorder(Recorder):
28    """Per-step CSV snapshots of every individual's lattice position +
29    a handful of attributes useful for color-coding the animation."""
30
31    def __init__(self, odir: pathlib.Path):
32        self.odir = odir / "lattice"
33        self.init_odir()
34
35    def write(self, population):
36        """
37        # OUTPUT SPECIFICATION
38        path: /lattice/step{step}.csv
39        filetype: csv
40        category: log
41        description: Per-individual lattice positions at a single step. One row per living individual: step, q, r, age, sex, lineage_id, ancestry_fraction. Sex and ancestry-related columns are -1 when their respective tracking is disabled.
42        trait granularity: individual
43        time granularity: snapshot
44        frequency parameter: LATTICE_RECORD_RATE
45        structure: CSV.
46        """
47        if not parametermanager.parameters.LATTICE_MODE:
48            return
49        if parametermanager.parameters.LATTICE_RECORD_RATE <= 0:
50            return
51        if population.positions is None or len(population) == 0:
52            return
53
54        step = variables.steps
55        should_skip = skip("LATTICE_RECORD_RATE")
56        is_first_step = step == 1
57        is_last_step = step == parametermanager.parameters.STEPS_PER_SIMULATION
58        if not (is_first_step or not should_skip or is_last_step):
59            return
60
61        n = len(population)
62        positions = population.positions  # (n, 2)
63        ages = population.ages  # (n,)
64        # sexes: existing AEGIS conventions vary; we just dump the raw int values
65        sexes = population.sexes if population.sexes is not None else np.full(n, -1, dtype=np.int32)
66        lineage_id = (
67            population.lineage_id if population.lineage_id is not None
68            else np.full(n, -1, dtype=np.int64)
69        )
70        if population.ancestry is not None:
71            ancestry = population.ancestry
72            # ancestry is bool, shape (n, ploidy, n_loci, bits_per_locus). Per-individual
73            # fraction = mean(True bits) across each individual's whole genome.
74            frac = ancestry.reshape(n, -1).mean(axis=1)
75        else:
76            frac = np.full(n, -1.0, dtype=np.float32)
77
78        path = self.odir / f"step{step}.csv"
79        with open(path, "w") as fh:
80            fh.write(HEADER)
81            for i in range(n):
82                fh.write(
83                    f"{step},{int(positions[i, 0])},{int(positions[i, 1])},"
84                    f"{int(ages[i])},{int(sexes[i])},{int(lineage_id[i])},{float(frac[i]):.4f}\n"
85                )
86        logging.debug(f"lattice snapshot recorded at step {step}.")