stress_operator/data/ood_split.py

"""Out-of-distribution (covariate-shift) split from void-geometry statistics (Gate D).

No new FEM is needed: we re-use the existing 2000 samples but split them by a per-sample
geometric statistic of the void boundary radii ``rr`` (shape ``(n, 42)``). The model trains on
the lower percentile (typical voids) and is tested on the upper percentile (large/extreme voids
it never saw) — a genuine covariate shift in geometry.

Split rule (documented exactly, master plan §2.6):
  stat_i = STAT(rr_i)                          # default STAT = max radius ("void size")
  threshold = quantile(stat, train_frac)        # default train_frac = 0.80
  in-distribution (train+id-test): stat_i <= threshold   (lower ~80%)
  out-of-distribution (ood-test) : stat_i >  threshold   (upper ~20%)
The in-distribution pool is further split into train / id-test by ``id_test_frac`` (held-out,
seeded) so in-dist and OOD errors are both measured on data unseen during training.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Dict

import numpy as np
import torch

STATS = {
    "max": lambda rr: rr.max(axis=1),       # largest radius -> void size (primary)
    "mean": lambda rr: rr.mean(axis=1),
    "std": lambda rr: rr.std(axis=1),        # radius spread -> "lobiness"
    "range": lambda rr: rr.max(axis=1) - rr.min(axis=1),
}


def geometry_stat(rr: torch.Tensor, kind: str = "max") -> np.ndarray:
    if kind not in STATS:
        raise ValueError(f"unknown stat {kind!r}; choose from {list(STATS)}")
    return STATS[kind](rr.detach().cpu().numpy())


@dataclass
class OODSplit:
    train_idx: np.ndarray       # in-distribution, used for training
    id_test_idx: np.ndarray     # in-distribution, held out for testing
    ood_test_idx: np.ndarray    # out-of-distribution (extreme geometry), held out
    threshold: float
    stat_kind: str
    info: Dict


def make_ood_split(
    rr: torch.Tensor,
    stat_kind: str = "max",
    train_frac: float = 0.80,
    id_test_frac: float = 0.20,
    seed: int = 0,
) -> OODSplit:
    """Build a covariate-shift split from ``rr`` statistics. Deterministic given ``seed``."""
    stat = geometry_stat(rr, stat_kind)            # (n,)
    n = stat.shape[0]
    threshold = float(np.quantile(stat, train_frac))

    in_dist = np.where(stat <= threshold)[0]
    ood = np.where(stat > threshold)[0]
    # A strict `> quantile` split yields an empty OOD group if the stat ties at its max
    # (degenerate distribution). Guard so `.min()/.max()` below cannot crash (audit bug 2);
    # real continuous `rr` data is non-degenerate (n_ood ~ 400 at train_frac=0.8).
    if ood.size == 0:
        raise ValueError(
            f"OOD split is empty: stat ({stat_kind}) has no values above the "
            f"{train_frac:.0%} quantile (likely tied at the maximum). Lower train_frac "
            f"or choose a different stat."
        )

    rng = np.random.default_rng(seed)
    perm = rng.permutation(in_dist)
    n_id_test = int(round(len(in_dist) * id_test_frac))
    id_test_idx = np.sort(perm[:n_id_test])
    train_idx = np.sort(perm[n_id_test:])

    info = {
        "n_total": int(n),
        "n_train": int(train_idx.size),
        "n_id_test": int(id_test_idx.size),
        "n_ood_test": int(ood.size),
        "threshold": threshold,
        "stat_kind": stat_kind,
        "stat_train_max": float(stat[train_idx].max()) if train_idx.size else float("nan"),
        "stat_ood_min": float(stat[ood].min()),
        "stat_ood_max": float(stat[ood].max()),
    }
    return OODSplit(
        train_idx=train_idx,
        id_test_idx=id_test_idx,
        ood_test_idx=np.sort(ood),
        threshold=threshold,
        stat_kind=stat_kind,
        info=info,
    )