MachineSex/src/neural/config.py

"""Resolved run configuration for a neural (Layer 1.5) lineage.

Mirrors the ``knowledge.config`` idiom exactly: frozen dataclasses with a ``from_dict``
that fills defaults and rejects unknown keys via ``knowledge.config._sub``. The grounding,
re-mint, and metrics blocks are *reused verbatim* from ``knowledge.config`` so the neural
runner speaks the same schema as Layer 1 (grounding ``m``, the ``g -> m`` conversion, the
re-mint gate, and the KL/support floors are all identical). Only the data source
(``synthetic``) and the model (``model``) are neural-specific.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field, replace
from typing import Any, Mapping

from knowledge.config import GroundingCfg, MetricsCfg, RemintCfg, _sub


@dataclass(frozen=True)
class SyntheticCfg:
    """The fully-synthetic mode-truth and observation grammar.

    The first seven fields are the Layer-1 ``TruthCfg`` knobs (they build ``p*`` over the
    ``K`` modes via ``knowledge.truth.make_true_distribution``). The remaining fields
    define how a mode is rendered to a categorical token sequence: an *identity* segment
    that encodes the mode losslessly (read by the exact oracle) followed by a *style*
    segment of within-mode stochastic tokens (so a real generative model has a
    distribution to learn, not just a lookup table).
    """

    K: int
    R: int = 1
    tail: str = "zipf"
    zipf_s: float = 1.1
    tail_frac: float = 0.5
    tail_threshold: float = 1e-3
    init: str = "uniform"        # initial p_0 over modes: {uniform, truth}
    style_len: int = 4           # style-segment length (within-mode entropy)
    style_vocab: int = 6         # style token alphabet size
    id_base: int = 2             # identity segment encodes the mode in this base

    @property
    def id_len(self) -> int:
        """Identity-segment length: fewest base-``id_base`` digits to index ``K`` modes."""
        if self.K <= 1:
            return 1
        return max(1, math.ceil(math.log(self.K, self.id_base)))

    @property
    def vocab(self) -> int:
        """Token alphabet size (shared by identity and style segments)."""
        return max(self.id_base, self.style_vocab)

    @property
    def seq_len(self) -> int:
        """Total observation length in tokens."""
        return self.id_len + self.style_len


@dataclass(frozen=True)
class MnistCfg:
    """Real-MNIST mode-truth: a Zipf ``p*`` over ``K = n_classes * style_bins`` modes.

    A mode is ``(digit class, stroke-thickness bin)`` under the fixed bijection
    ``mode = class * style_bins + bin``. The first seven fields are the Layer-1 ``TruthCfg``
    knobs (they build the Zipf ``p*`` over the ``K`` modes via ``make_true_distribution``);
    the rest govern the image tier. Unlike ``SyntheticCfg`` there is no rendering grammar —
    observations are real images and the oracle is a frozen classifier.
    """

    K: int = 30
    R: int = 1
    tail: str = "zipf"
    zipf_s: float = 1.5          # steep enough that the rarest ~18/30 modes form a real tail
    tail_frac: float = 0.5
    tail_threshold: float = 1e-2  # modes with p* < 0.01 are "tail" (~9% of the mass)
    init: str = "truth"          # initial p_0 over modes: {uniform, truth}
    n_classes: int = 10          # MNIST digit classes
    style_bins: int = 3          # S: per-class stroke-thickness quantile bins (K = n_classes*S)
    data_root: str = "data"      # gitignored MNIST download dir

    def __post_init__(self) -> None:
        if self.K != self.n_classes * self.style_bins:
            raise ValueError(
                f"K ({self.K}) must equal n_classes*style_bins "
                f"({self.n_classes}*{self.style_bins}={self.n_classes * self.style_bins})")


@dataclass(frozen=True)
class OracleCfg:
    """Frozen-classifier oracle training/caching (MNIST tier)."""

    epochs: int = 5
    lr: float = 1.0e-3
    batch_size: int = 256
    cache: str = "models/mnist_cnn.pt"   # gitignored checkpoint; its hash goes in the manifest


@dataclass(frozen=True)
class ModelCfg:
    """The generative learner. ``kind`` selects the architecture behind a thin adapter.

    Neural hyperparameters are ignored by the ``histogram`` bridge model.
    """

    kind: str = "histogram"      # {histogram, rnn, vae, mlp, convvae}
    hidden: int = 64
    embed: int = 16
    epochs: int = 30
    lr: float = 1.0e-3
    batch_size: int = 256
    device: str = "auto"         # {auto, cpu, cuda}
    latent: int = 16             # VAE latent dimension (VAE only)
    beta: float = 1.0            # VAE KL weight (VAE only)
    # Samples used to estimate a neural model's mode distribution by generate-and-classify
    # (ignored by the exact histogram bridge). Larger -> less measurement noise on p_hat.
    n_eval: int = 8000


@dataclass(frozen=True)
class NeuralDynamicsCfg:
    """Generational dynamics: drift strength ``n`` + reused grounding/re-mint blocks."""

    n: int = 4000                # pupil training-sample size (drift strength ~ 1/n)
    grounding: GroundingCfg = field(default_factory=GroundingCfg)
    remint: RemintCfg = field(default_factory=RemintCfg)


@dataclass(frozen=True)
class NeuralLineageCfg:
    """A fully-resolved neural lineage configuration."""

    synthetic: SyntheticCfg
    model: ModelCfg = field(default_factory=ModelCfg)
    dynamics: NeuralDynamicsCfg = field(default_factory=NeuralDynamicsCfg)
    generations: int = 30
    metrics: MetricsCfg = field(default_factory=MetricsCfg)

    @staticmethod
    def from_dict(cfg: Mapping[str, Any]) -> "NeuralLineageCfg":
        """Build a validated NeuralLineageCfg from a nested mapping, filling defaults."""
        if isinstance(cfg, NeuralLineageCfg):
            return cfg
        synthetic = _sub(cfg.get("synthetic", {}), SyntheticCfg)
        model = _sub(cfg.get("model", {}), ModelCfg)
        dyn_raw = dict(cfg.get("dynamics", {}))
        dynamics = NeuralDynamicsCfg(
            n=dyn_raw.get("n", NeuralDynamicsCfg.n),
            grounding=_sub(dyn_raw.get("grounding", {}), GroundingCfg),
            remint=_sub(dyn_raw.get("remint", {}), RemintCfg),
        )
        metrics = _sub(cfg.get("metrics", {}), MetricsCfg)
        return NeuralLineageCfg(
            synthetic=synthetic,
            model=model,
            dynamics=dynamics,
            generations=int(cfg.get("generations", NeuralLineageCfg.generations)),
            metrics=metrics,
        )

    def replace(self, **kw) -> "NeuralLineageCfg":
        return replace(self, **kw)