MachineSex/src/inheritance/speciation.py

"""E12 — model speciation: when two diverged models are too incompatible to recombine (merge).

The evolution-of-sex frame (E7–E11) says recombining complementary parents beats copying. This module
asks the geneticist's question underneath it: **how far can two lineages diverge before recombination
stops working?** In biology the answer is *reproductive isolation* via **Bateson–Dobzhansky–Muller
incompatibilities** (BDMIs): alleles that are each benign on their own lineage's background but
deleterious *in combination*, so a recombinant (a hybrid) inherits combinations selection never tested.
A merged model is a single recombinant genotype — an F2-like *hybrid breakdown / recombination-load*
object, not an F1 — so the predicted signature as divergence grows is
**compatible → outbreeding depression → hybrid inviability**, arriving *earlier* the more epistatic the
capability landscape. This is the analytic core of the paper's speciation claim; it is confirmed in
sign, not magnitude, by real weights elsewhere.

Two landscapes, one runner:

- ``bdm`` — the controllable, canonical construction. An ancestor; two lineages each substitute a
  *disjoint* set of loci (so each parent is adapted and neither carries an incompatibility); a fraction
  ``rho`` of the cross-lineage locus pairs are BDMIs with penalty ``s``. Sweeping the divergence ``d``
  (total substitutions) yields the fitness curve and the **Orr–Turelli snowball** — the number of
  incompatibilities grows ~``(d/2)^2``, so hybrid fitness falls *super-linearly*.
- ``nk`` — the emergent version. Parents are *local optima* (hill-climbed) on a Kauffman NK landscape;
  recombining them exposes broken co-adapted blocks. Sweeping the ruggedness ``K`` at matched divergence
  isolates the paper's wedge: **at equal divergence, mergeability is governed by epistasis**, which no
  divergence-only ML predictor captures.

Pure seeded NumPy on the existing genotype machinery — bitwise-reproducible from one master seed, no
external simulator (whose separate, version-unstable RNG would break that guarantee).
"""

from __future__ import annotations

from typing import Any, Mapping

import numpy as np
import pandas as pd

from .genotype import bits_to_index, crossover, genotype_bits, hill_climb, nk_fitness
from .seeding import spawn_seeds


# --------------------------------------------------------------------------- BDM construction
def _bdm_point(L: int, d: int, rho: float, s: float, beta: float, rate: float,
               n_off: int, rng: np.random.Generator) -> dict:
    """One divergence point of the BDM model (one replicate's random locus assignment + offspring)."""
    half = min(d, L - (L % 2)) // 2
    loci = rng.permutation(L)
    S1, S2 = loci[:half], loci[half:2 * half]                  # disjoint substituted loci per lineage
    p1 = np.zeros(L, dtype=np.int8); p1[S1] = 1                 # parent 1: derived at S1
    p2 = np.zeros(L, dtype=np.int8); p2[S2] = 1                 # parent 2: derived at S2
    dmi = [(int(a), int(b)) for a in S1 for b in S2 if rng.random() < rho]  # cross-lineage BDMIs

    def fitness(bits: np.ndarray) -> np.ndarray:
        b = np.atleast_2d(bits)
        val = beta * b.sum(1).astype(float)                    # additive: each derived allele is adaptive
        for a, bl in dmi:                                      # BDMI: penalty only if BOTH derived alleles present
            val -= s * ((b[:, a] == 1) & (b[:, bl] == 1))
        return val

    f1, f2 = float(fitness(p1)[0]), float(fitness(p2)[0])      # parents carry no incompatibility (disjoint)
    parents = np.stack([p1, p2])
    offs = np.stack([crossover(parents, rate, rng) for _ in range(n_off)])
    fo = fitness(offs)
    worse_parent = min(f1, f2)
    realized = np.zeros(n_off)
    for a, bl in dmi:
        realized += (offs[:, a] == 1) & (offs[:, bl] == 1)
    return {"divergence": int(2 * half), "n_dmi": len(dmi),
            "parent_fitness": (f1 + f2) / 2.0,
            "offspring_fitness": float(fo.mean()),
            "outbreeding_depression": float((fo < worse_parent).mean()),   # hybrid worse than either parent
            "isolation": float((fo < 0.0).mean()),                          # hybrid below the ancestor = inviable
            "incompatibilities": float(realized.mean())}


# --------------------------------------------------------------------------- NK construction
def _nk_point(L: int, K: int, landscape_seed: int, rate: float, n_pairs: int,
              n_off: int, rng: np.random.Generator) -> dict:
    """Aggregate over ``n_pairs`` random parent-pairs (local optima) on one NK landscape at ruggedness K."""
    F = nk_fitness(L, K, landscape_seed)
    bits = genotype_bits(L)
    divs, gaps, od = [], [], []
    for _ in range(n_pairs):
        g1 = hill_climb(F, L, int(rng.integers(1 << L)))
        g2 = hill_climb(F, L, int(rng.integers(1 << L)))
        b1, b2 = bits[g1], bits[g2]
        parents = np.stack([b1, b2])
        offs = [crossover(parents, rate, rng) for _ in range(n_off)]
        fo = np.array([F[bits_to_index(o)] for o in offs])
        worse = min(F[g1], F[g2])
        divs.append(int((b1 != b2).sum()))
        gaps.append(float(fo.mean() - worse))                  # >0: recombination helps; <0: it hurts
        od.append(float((fo < worse).mean()))
    return {"K": K, "divergence": float(np.mean(divs)),
            "offspring_minus_parent": float(np.mean(gaps)),
            "outbreeding_depression": float(np.mean(od))}


# --------------------------------------------------------------------------- runner
def run_speciation(cfg: Mapping[str, Any], seed: int) -> pd.DataFrame:
    """Run the speciation sweep (``landscape`` = ``bdm`` or ``nk``); return a tidy per-point DataFrame.

    Args:
        cfg (Mapping): resolved config with a ``speciation`` block.
        seed (int): master seed (all randomness derives from it via ``spawn_seeds``).

    Returns:
        pd.DataFrame: one row per swept point x replicate, with the metric columns.
    """
    spec = dict(cfg["speciation"])
    landscape = spec.get("landscape", "bdm")
    L = int(spec.get("L", 16))
    rate = float(spec.get("recomb_rate", 0.5))
    n_off = int(spec.get("n_offspring", 400))
    reps = int(cfg.get("n_replicates", spec.get("reps", 20)))
    rows: list[dict] = []
    seeds = spawn_seeds(seed, reps)

    if landscape == "bdm":
        rhos = spec.get("rho", [0.1, 0.25, 0.5])
        rhos = rhos if isinstance(rhos, (list, tuple)) else [rhos]
        divergences = spec.get("divergences", list(range(0, L + 1, 2)))
        s, beta = float(spec.get("s", 1.0)), float(spec.get("beta", 1.0))
        for rep, ss in enumerate(seeds):
            rng = np.random.default_rng(int(ss.generate_state(1)[0]))
            for rho in rhos:
                for d in divergences:
                    row = _bdm_point(L, int(d), float(rho), s, beta, rate, n_off, rng)
                    row.update({"landscape": "bdm", "rho": float(rho), "replicate": rep})
                    rows.append(row)
    elif landscape == "nk":
        Ks = spec.get("K", [0, 2, 4, 6, 8])
        Ks = Ks if isinstance(Ks, (list, tuple)) else [Ks]
        n_pairs = int(spec.get("n_pairs", 40))
        for rep, ss in enumerate(seeds):
            state = ss.generate_state(2)
            rng = np.random.default_rng(int(state[0]))
            for K in Ks:
                row = _nk_point(L, int(K), int(state[1]), rate, n_pairs, n_off, rng)
                row.update({"landscape": "nk", "replicate": rep})
                rows.append(row)
    else:
        raise ValueError(f"unknown speciation landscape {landscape!r} (expected bdm|nk)")

    return pd.DataFrame(rows)