MachineSex/src/knowledge/lineage.py

"""Run a single lineage for T generations (blueprint 2.7 interface).

``run_lineage`` is the core driver every experiment builds on. It accepts a plain nested
mapping (the scientific-validation conformance tests pass one directly) or a
:class:`~knowledge.config.LineageCfg`, and returns a tidy DataFrame with one row per
generation (0..T inclusive) carrying every blueprint-2.3 metric, global and per-region.
"""

from __future__ import annotations

from typing import Any, Mapping

import numpy as np
import pandas as pd

from .config import LineageCfg
from .metrics import (forward_kl, heterozygosity, per_region, support_size,
                      tail_band_metrics, tail_mass)
from .step import StepCtx, allocate_m, generation_step
from .truth import make_true_distribution, uniform_init

N_BANDS = 4  # rarity bands for per-band tail-survival logging (blueprint pred. 4)


def run_lineage(cfg: Mapping[str, Any] | LineageCfg, seed: int) -> pd.DataFrame:
    """Run one lineage and return per-generation metrics.

    Args:
        cfg (Mapping | LineageCfg): Resolved lineage configuration (blueprint 2.7 schema).
        seed (int): Seed for this replicate; the run is a pure function of (cfg, seed).

    Returns:
        pd.DataFrame: One row per generation 0..T with columns ``generation``,
        ``heterozygosity``, ``forward_kl``, ``tail_mass``, ``support_size``, and
        per-region ``H_region_{r}`` / ``tail_region_{r}`` columns.
    """
    cfg = LineageCfg.from_dict(cfg)
    td = make_true_distribution(
        cfg.truth.K, cfg.truth.R, cfg.truth.tail, cfg.truth.tail_frac,
        cfg.truth.zipf_s, seed, tail_threshold=cfg.truth.tail_threshold,
    )
    p_star_orig = td.p_star          # the ORIGINAL truth; forward_kl is always vs this
    regions = td.regions
    tail_mask = td.tail_mask         # tail defined on the original truth
    kl_floor = cfg.metrics.kl_floor
    eps = cfg.metrics.support_eps

    rng = np.random.default_rng(seed)

    if cfg.truth.init == "uniform":
        p = uniform_init(cfg.truth.K)
    elif cfg.truth.init == "truth":
        p = p_star_orig.copy()
    else:
        raise ValueError(f"unknown init {cfg.truth.init!r} (expected uniform|truth)")

    p_star_eff = p_star_orig.copy()  # grounding reference; may be re-minted (E6)
    exercised = cfg.dynamics.grounding.exercised
    exercised = np.asarray(exercised) if exercised is not None else None
    m_vector = allocate_m(cfg.dynamics.grounding.m, cfg.truth.R,
                          cfg.dynamics.grounding.policy, exercised)
    step_ctx = StepCtx(
        n=cfg.dynamics.n,
        m_vector=m_vector,
        policy=cfg.dynamics.grounding.policy,
        regions=regions,
        selection_mode=cfg.dynamics.selection.mode,
        novelty_alpha=cfg.dynamics.selection.novelty_alpha,
        kernel=cfg.dynamics.kernel,
    )
    remint = cfg.dynamics.remint

    rows: list[dict] = []

    head_mask = ~tail_mask
    n_tail = int(tail_mask.sum())
    n_head = int(head_mask.sum())

    def record(t: int, p: np.ndarray) -> None:
        # Tail-vs-head SURVIVAL is the honest collapse signature: under neutral drift the
        # mean tail *mass* is a martingale (conserved), but tail *items* go extinct first.
        row = {
            "generation": t,
            "heterozygosity": heterozygosity(p),
            "forward_kl": forward_kl(p_star_orig, p, kl_floor),
            "tail_mass": tail_mass(p, tail_mask),
            "support_size": support_size(p, eps),
            "tail_support": int(np.sum(p[tail_mask] > eps)),
            "head_support": int(np.sum(p[head_mask] > eps)),
            "tail_frac_alive": (float(np.mean(p[tail_mask] > eps)) if n_tail else 0.0),
            "head_frac_alive": (float(np.mean(p[head_mask] > eps)) if n_head else 0.0),
            # Truth-mass-weighted tail coverage: share of the tail's TRUE mass carried by
            # still-alive items. Bounded, monotone in grounding, and the honest "mass
            # rescued" companion to tail_frac_alive (raw tail_mass is a drift martingale).
            "tail_truth_mass_alive": (
                float(p_star_orig[tail_mask][p[tail_mask] > eps].sum()
                      / p_star_orig[tail_mask].sum()) if n_tail else 0.0),
        }
        if n_tail >= N_BANDS:  # per-rarity-band survival (band 0 = rarest); see metrics
            fa, _ = tail_band_metrics(p, p_star_orig, tail_mask, n_bands=N_BANDS,
                                      alive_eps=eps)
            for b in range(N_BANDS):
                row[f"band{b}_alive"] = fa[b]
        if cfg.truth.R > 1:  # per-region columns only when there is >1 region
            for r, v in per_region(heterozygosity, p, regions).items():
                row[f"H_region_{r}"] = v
            for r, v in per_region(tail_mass, p, regions, tail_mask).items():
                row[f"tail_region_{r}"] = v
            # per-region tail-item survival (E3's honest metric: how much of each
            # region's rare tail is kept alive, not just its martingale mass)
            for r in range(cfg.truth.R):
                region_tail = (regions == r) & tail_mask
                row[f"tailalive_region_{r}"] = (
                    float(np.mean(p[region_tail] > eps)) if region_tail.any() else 0.0)
        rows.append(row)

    record(0, p)
    for t in range(1, cfg.generations + 1):
        p = generation_step([p], p_star_eff, step_ctx, rng)
        if (remint.enabled and remint.period and t % remint.period == 0):
            # Re-mint (founder event): the current distribution becomes the new grounding
            # reference and the original truth is discarded for grounding purposes. Gated
            # on diversity: only re-mint if H is high enough (E6). forward_kl stays vs the
            # original truth, so a collapsed re-mint locks KL high forever.
            if remint.H_gate is None or heterozygosity(p) >= remint.H_gate:
                p_star_eff = p.copy()
        record(t, p)

    return pd.DataFrame(rows)