aegis_sim.submodels.lattice

  1"""Hexagonal lattice spatial model (opt-in via LATTICE_MODE).
  2
  3When enabled, individuals occupy positions on a toroidal hexagonal lattice
  4(parallelogram with wraparound on both axes, axial coordinates (q, r)).
  5The lattice constrains:
  6
  7  - mating: females search expanding rings outward for a fertile male
  8  - offspring placement: random adjacent empty cell, birth fails if none
  9  - migration: per-step probability of moving to a random adjacent empty
 10               cell, plus a rare long-distance dispersal to any empty cell
 11
 12When LATTICE_MODE is False (default), this submodel is not instantiated and
 13the simulation runs exactly as before. All entry points are no-ops on the
 14None-positions path; nothing in the bioreactor changes unless the caller
 15checks LATTICE_MODE first.
 16
 17Hexagonal geometry: axial coordinates (q, r) with q in [0, ROWS),
 18r in [0, COLS), toroidal wraparound on both axes. The six neighbors of
 19(q, r) are:
 20
 21    ( q+1, r   ), ( q-1, r   ),
 22    ( q,   r+1 ), ( q,   r-1 ),
 23    ( q+1, r-1 ), ( q-1, r+1 )
 24
 25The "ring" of distance d around (q, r) contains 6*d cells (for d >= 1).
 26Distance 1 is the 6 immediate neighbors; distance 2 is the next 12 cells;
 27distance 3 is 18; etc.
 28"""
 29
 30import logging
 31from typing import Optional, Tuple
 32
 33import numpy as np
 34
 35
 36# Module-level singleton state. Populated by init() when LATTICE_MODE is on.
 37# occupancy is a 2D bool array — True = cell is occupied. Population.positions
 38# is the source of truth for *which* individual is at a cell; the lattice
 39# only tracks occupancy/vacancy. This avoids correctness bugs when the
 40# population is reindexed by mortality or merges.
 41_state = {
 42    "rows": 0,
 43    "cols": 0,
 44    "occupancy": None,  # 2D bool array; True = occupied
 45    "rng": None,
 46}
 47
 48
 49def init(LATTICE_MODE, INITIAL_POPULATION_SIZE, LATTICE_TARGET_DENSITY,
 50         RESOURCE_MAXIMUM_AMOUNT, rng_seed=None):
 51    """Initialise the lattice singleton.
 52
 53    Lattice size is computed from the expected carrying capacity (we use
 54    RESOURCE_MAXIMUM_AMOUNT as a proxy when available, falling back to
 55    INITIAL_POPULATION_SIZE) and LATTICE_TARGET_DENSITY:
 56
 57        n_cells = expected_carrying_capacity / LATTICE_TARGET_DENSITY
 58
 59    The lattice is then sized as a roughly square parallelogram:
 60    rows = ceil(sqrt(n_cells)), cols = ceil(n_cells / rows).
 61    """
 62    if not LATTICE_MODE:
 63        # Spatial model disabled. Keep state empty; callers must guard on LATTICE_MODE.
 64        for k in _state:
 65            _state[k] = None if k != "rows" and k != "cols" else 0
 66        return
 67
 68    expected_capacity = max(int(RESOURCE_MAXIMUM_AMOUNT or 0), int(INITIAL_POPULATION_SIZE))
 69    n_cells_target = max(1, int(np.ceil(expected_capacity / max(LATTICE_TARGET_DENSITY, 1e-6))))
 70    rows = max(1, int(np.ceil(np.sqrt(n_cells_target))))
 71    cols = max(1, int(np.ceil(n_cells_target / rows)))
 72
 73    _state["rows"] = rows
 74    _state["cols"] = cols
 75    _state["occupancy"] = np.zeros((rows, cols), dtype=bool)
 76    _state["rng"] = np.random.default_rng(rng_seed)
 77
 78    logging.info(
 79        "Lattice initialised: %d rows x %d cols = %d cells (target density %.2f, expected capacity %d)",
 80        rows, cols, rows * cols, LATTICE_TARGET_DENSITY, expected_capacity,
 81    )
 82
 83
 84def _wrap(q, r) -> Tuple[int, int]:
 85    """Toroidal wraparound to canonical (q, r)."""
 86    return int(q) % _state["rows"], int(r) % _state["cols"]
 87
 88
 89# The six axial-coordinate offsets to neighbouring hex cells.
 90_NEIGHBOUR_OFFSETS = np.array(
 91    [(1, 0), (-1, 0), (0, 1), (0, -1), (1, -1), (-1, 1)],
 92    dtype=np.int32,
 93)
 94
 95
 96def neighbours(q: int, r: int) -> np.ndarray:
 97    """Return the (q, r) coordinates of the 6 immediate neighbours of cell (q, r).
 98    Wrapped onto the torus. Shape: (6, 2)."""
 99    offsets = _NEIGHBOUR_OFFSETS
100    qs = (q + offsets[:, 0]) % _state["rows"]
101    rs = (r + offsets[:, 1]) % _state["cols"]
102    return np.stack([qs, rs], axis=1)
103
104
105def ring(q: int, r: int, radius: int) -> np.ndarray:
106    """Return all (q, r) cells at hex-distance exactly `radius` from (q, r).
107
108    The ring at distance d has 6*d cells (d >= 1). Wrapped onto the torus.
109    Shape: (6 * radius, 2). For radius == 0, returns just (q, r).
110
111    Algorithm: start at corner (q + radius, r - radius) (radius steps in
112    direction +1,-1) and walk `radius` cells in each of the six side
113    directions in turn. The walking directions form a CW cycle around the
114    hex; see _RING_WALK_DIRS below for the order verified against the
115    six immediate neighbours at radius=1.
116    """
117    if radius == 0:
118        return np.array([[q % _state["rows"], r % _state["cols"]]], dtype=np.int32)
119
120    cur_q = q + radius
121    cur_r = r - radius
122    out = np.empty((6 * radius, 2), dtype=np.int32)
123    idx = 0
124    for dq, dr in _RING_WALK_DIRS:
125        for _ in range(radius):
126            out[idx, 0] = cur_q % _state["rows"]
127            out[idx, 1] = cur_r % _state["cols"]
128            idx += 1
129            cur_q += dq
130            cur_r += dr
131
132    return out
133
134
135# Walking directions around a hex ring, in CW order starting from the
136# corner at (q + radius, r - radius). Verified against immediate neighbours.
137_RING_WALK_DIRS = (
138    (-1, 0),   # west
139    (-1, 1),   # south-west
140    (0, 1),    # south
141    (1, 0),    # east
142    (1, -1),   # north-east
143    (0, -1),   # north
144)
145
146
147def is_empty(q: int, r: int) -> bool:
148    q, r = _wrap(q, r)
149    return not _state["occupancy"][q, r]
150
151
152def claim(q: int, r: int) -> None:
153    """Mark cell (q, r) as occupied. Raises if already occupied — caller's
154    responsibility to vacate first when moving an individual."""
155    q, r = _wrap(q, r)
156    if _state["occupancy"][q, r]:
157        raise RuntimeError(f"Cannot claim cell ({q}, {r}): already occupied")
158    _state["occupancy"][q, r] = True
159
160
161def vacate(q: int, r: int) -> None:
162    q, r = _wrap(q, r)
163    _state["occupancy"][q, r] = False
164
165
166def random_empty_anywhere() -> Optional[Tuple[int, int]]:
167    """Pick a uniformly-random empty cell from the entire lattice. None if full."""
168    empties = np.argwhere(~_state["occupancy"])
169    if len(empties) == 0:
170        return None
171    pick = _state["rng"].integers(0, len(empties))
172    return int(empties[pick, 0]), int(empties[pick, 1])
173
174
175def random_empty_adjacent(q: int, r: int) -> Optional[Tuple[int, int]]:
176    """Pick a uniformly-random empty cell from the 6 neighbours of (q, r).
177    None if all neighbours are occupied."""
178    cells = neighbours(q, r)
179    empties = [(int(c[0]), int(c[1])) for c in cells if not _state["occupancy"][c[0], c[1]]]
180    if not empties:
181        return None
182    return empties[_state["rng"].integers(0, len(empties))]
183
184
185def resync_occupancy_from_positions(*position_arrays) -> None:
186    """Rebuild the occupancy grid from one or more positions arrays.
187
188    Call after any operation that mutates the population or the egg
189    pool (kills, hatching, reproduction with delayed hatching, merges)
190    so the lattice's occupancy stays consistent with the source of truth.
191
192    Multiple arrays may be passed when eggs occupy lattice cells during
193    incubation: pass both Population.positions and Eggs.positions.
194    Sentinel rows (any coordinate < 0) are skipped — these represent
195    individuals not yet placed on the lattice. Cheap: O(n_cells) clear
196    + O(total positions) set.
197    """
198    if _state["occupancy"] is None:
199        return
200    _state["occupancy"].fill(False)
201    for positions in position_arrays:
202        if positions is None or len(positions) == 0:
203            continue
204        # Skip sentinel rows (e.g. (-1, -1) for eggs that failed placement)
205        valid = (positions[:, 0] >= 0) & (positions[:, 1] >= 0)
206        if not valid.any():
207            continue
208        qs = positions[valid, 0] % _state["rows"]
209        rs = positions[valid, 1] % _state["cols"]
210        _state["occupancy"][qs, rs] = True
211
212
213def migrate(positions: np.ndarray, migration_rate: float, migration_long_rate: float) -> None:
214    """Migrate each individual on the lattice in-place.
215
216    For each individual, roll a single random number:
217      < migration_long_rate              -> long-distance jump to a random
218                                            empty cell anywhere on the lattice
219      < migration_long_rate + migration_rate -> local move to a random adjacent
220                                            empty cell
221      otherwise                          -> stay put
222
223    Movement only happens if a target cell is available; otherwise the
224    individual stays in place. `positions` is mutated in place; the
225    occupancy grid is kept in sync.
226
227    Long-range first so the two probabilities are interpretable
228    independently (each has its stated meaning regardless of the other).
229    """
230    if _state["occupancy"] is None or positions is None or len(positions) == 0:
231        return
232    if migration_rate <= 0 and migration_long_rate <= 0:
233        return
234
235    rng = _state["rng"]
236    n = len(positions)
237    rolls = rng.random(n)
238
239    # Process long-range first so a successful long jump precludes a local move.
240    long_threshold = migration_long_rate
241    local_threshold = migration_long_rate + migration_rate
242
243    for i in range(n):
244        roll = rolls[i]
245        if roll >= local_threshold:
246            continue  # neither move
247        cur_q = int(positions[i, 0])
248        cur_r = int(positions[i, 1])
249        if roll < long_threshold:
250            target = random_empty_anywhere()
251        else:
252            target = random_empty_adjacent(cur_q, cur_r)
253        if target is None:
254            continue  # no destination, stay put
255        # Move: vacate old cell, claim new cell, update positions
256        vacate(cur_q, cur_r)
257        claim(target[0], target[1])
258        positions[i, 0] = target[0]
259        positions[i, 1] = target[1]
260
261
262def n_empty() -> int:
263    """Count of empty cells on the lattice. Diagnostic; not hot-path."""
264    if _state["occupancy"] is None:
265        return 0
266    return int((~_state["occupancy"]).sum())
267
268
269def assign_initial_positions(n: int) -> np.ndarray:
270    """Assign n unique empty cells to the initial population. Returns an
271    (n, 2) int32 array of (q, r) coordinates. Cells are marked occupied.
272    """
273    if _state["occupancy"] is None:
274        raise RuntimeError("Lattice not initialised; call submodels.lattice.init(...) first")
275
276    rows, cols = _state["rows"], _state["cols"]
277    n_cells = rows * cols
278    if n > n_cells:
279        raise ValueError(
280            f"Cannot place {n} individuals on a lattice of {n_cells} cells "
281            f"(rows={rows}, cols={cols}). Increase LATTICE_TARGET_DENSITY or "
282            f"reduce INITIAL_POPULATION_SIZE."
283        )
284
285    empties = np.argwhere(~_state["occupancy"])
286    pick = _state["rng"].choice(len(empties), size=n, replace=False)
287    chosen = empties[pick].astype(np.int32)
288    _state["occupancy"][chosen[:, 0], chosen[:, 1]] = True
289    return chosen