sensorflow/hardening/sampling.py

"""Corrected scenario ranking and stratified sampling with reweighting.

Fixes the audit's sampling findings for the legacy paths (F-020: arbitrary
rarity coefficients and confidence-from-rarity; "most interesting first"
ordering biasing P(failure) estimates):

1. Ranking is deterministic and evidence-decomposed (no unseeded randomness,
   no collapsed opaque score without its components).
2. Evaluation sets are drawn by STRATIFIED sampling with explicit quotas, and
   every drawn item carries an inverse-inclusion-probability weight so
   fleet-level rates can be estimated without oversampling bias
   (Horvitz-Thompson, same estimator megaeval.sampling already uses — that
   implementation was verified correct, F-029).
3. Novelty is real geometry over embeddings: kNN distance and centroid
   distance, with the trade-offs documented on each scorer.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np


class Stratum(str, Enum):
    REPRESENTATIVE = "Representative"
    NOMINAL = "Nominal"
    RARE = "Rare"
    SAFETY_CRITICAL = "SafetyCritical"
    KNOWN_FAILURES = "KnownFailures"
    NOVEL = "Novel"
    DISTRIBUTION_SHIFTED = "DistributionShifted"


#: Default sampling quotas (fraction of the evaluation budget per stratum).
#: These are POLICY parameters — explicit, versioned, and overridable — not
#: buried magic numbers. They intentionally oversample rare/safety strata;
#: the HT weights below undo that oversampling for population estimates.
DEFAULT_QUOTAS: Dict[Stratum, float] = {
    Stratum.REPRESENTATIVE: 0.30,
    Stratum.NOMINAL: 0.15,
    Stratum.RARE: 0.15,
    Stratum.SAFETY_CRITICAL: 0.15,
    Stratum.KNOWN_FAILURES: 0.10,
    Stratum.NOVEL: 0.10,
    Stratum.DISTRIBUTION_SHIFTED: 0.05,
}


@dataclass
class SampledItem:
    item_id: str
    stratum: Stratum
    inclusion_probability: float
    weight: float                # 1 / inclusion_probability (Horvitz-Thompson)


@dataclass
class StratifiedSample:
    items: List[SampledItem]
    quotas: Dict[Stratum, int]
    population_sizes: Dict[Stratum, int]
    seed: int

    def estimate_population_rate(self, indicator: Dict[str, float]) -> Optional[float]:
        """HT estimate of a population mean/rate from the stratified sample.

        indicator: item_id -> 0/1 (or real-valued) outcome measured on the
        sample. Returns sum(w_i * y_i) / sum over population size — i.e. the
        stratum-weighted mean that recovers the true population rate under
        oversampling (each stratum's sample mean weighted by its true share).
        """
        total_pop = sum(self.population_sizes.values())
        if total_pop == 0:
            return None
        est = 0.0
        for stratum, pop in self.population_sizes.items():
            in_stratum = [it for it in self.items if it.stratum == stratum
                          and it.item_id in indicator]
            if not in_stratum:
                if pop > 0:
                    return None  # a non-empty stratum with no measurements: refuse
                continue
            stratum_mean = float(np.mean([indicator[it.item_id] for it in in_stratum]))
            est += (pop / total_pop) * stratum_mean
        return est


def stratified_sample(
    population: Dict[Stratum, Sequence[str]],
    budget: int,
    quotas: Optional[Dict[Stratum, float]] = None,
    seed: int = 0,
) -> StratifiedSample:
    """Draw a stratified sample with explicit quotas and HT weights.

    - Quota fractions are converted to integer counts by largest-remainder
      allocation (fixes the naive-rounding caveat noted in the audit, F-029).
    - Sampling within a stratum is uniform without replacement, from a seeded
      generator: same seed -> same sample, always.
    - If a stratum has fewer items than its quota, the whole stratum is taken
      (inclusion probability 1) and the leftover budget is redistributed to
      the remaining strata proportionally.
    """
    quotas = quotas or DEFAULT_QUOTAS
    rng = np.random.default_rng(seed)

    # Largest-remainder integer allocation of the budget.
    fracs = {s: quotas.get(s, 0.0) for s in population}
    total_frac = sum(fracs.values()) or 1.0
    raw = {s: budget * f / total_frac for s, f in fracs.items()}
    alloc = {s: int(v) for s, v in raw.items()}
    remainders = sorted(raw, key=lambda s: raw[s] - alloc[s], reverse=True)
    for s in remainders[: budget - sum(alloc.values())]:
        alloc[s] += 1

    # Cap at stratum size; redistribute leftovers deterministically.
    leftover = 0
    for s in sorted(alloc, key=lambda s: s.value):
        n_avail = len(population[s])
        if alloc[s] > n_avail:
            leftover += alloc[s] - n_avail
            alloc[s] = n_avail
    while leftover > 0:
        expandable = [s for s in sorted(alloc, key=lambda s: s.value)
                      if alloc[s] < len(population[s])]
        if not expandable:
            break
        for s in expandable:
            if leftover == 0:
                break
            alloc[s] += 1
            leftover -= 1

    items: List[SampledItem] = []
    for s in sorted(population, key=lambda s: s.value):
        ids = sorted(population[s])  # order-independence: canonical order first
        k = alloc.get(s, 0)
        if k == 0 or not ids:
            continue
        chosen = rng.choice(len(ids), size=k, replace=False)
        p_inc = k / len(ids)
        for idx in sorted(chosen.tolist()):
            items.append(SampledItem(
                item_id=ids[idx], stratum=s,
                inclusion_probability=p_inc, weight=1.0 / p_inc,
            ))

    return StratifiedSample(
        items=items,
        quotas=alloc,
        population_sizes={s: len(v) for s, v in population.items()},
        seed=seed,
    )


# ------------------------------------------------------------------ novelty


def knn_novelty(embeddings: np.ndarray, reference: np.ndarray, k: int = 10) -> np.ndarray:
    """Novelty = mean Euclidean distance to the k nearest reference embeddings.

    Trade-offs: sensitive to LOCAL density, so it finds points in sparse
    pockets even inside the reference hull (good for multi-modal fleets);
    cost is O(n_query * n_reference) without an index — use the VectorDB
    interface (interfaces.py) with an ANN index at production scale. k too
    small -> noisy; k too large -> smooths away small novel clusters.
    """
    if len(reference) == 0:
        return np.full(len(embeddings), np.inf)
    k = min(k, len(reference))
    # Chunked exact distances; replace with ANN via VectorDB in production.
    out = np.empty(len(embeddings))
    for i in range(0, len(embeddings), 1024):
        chunk = embeddings[i:i + 1024]
        d = np.linalg.norm(chunk[:, None, :] - reference[None, :, :], axis=2)
        d.sort(axis=1)
        out[i:i + 1024] = d[:, :k].mean(axis=1)
    return out


def centroid_novelty(embeddings: np.ndarray, reference: np.ndarray) -> np.ndarray:
    """Novelty = distance to the reference centroid, scaled by mean spread.

    Trade-offs: O(n) and index-free, but UNIMODAL — it misses novel points
    that sit near the global centroid of a multi-modal distribution and
    over-flags legitimate points on the far side of a wide mode. Use as a
    cheap first pass; confirm with knn_novelty before routing on it.
    """
    if len(reference) == 0:
        return np.full(len(embeddings), np.inf)
    centroid = reference.mean(axis=0)
    spread = float(np.mean(np.linalg.norm(reference - centroid, axis=1))) or 1.0
    return np.linalg.norm(embeddings - centroid, axis=1) / spread


# ------------------------------------------------------------------ ranking


@dataclass
class RankedScenario:
    scenario_id: str
    score: float
    components: Dict[str, float] = field(default_factory=dict)


def rank_scenarios(
    scenario_ids: Sequence[str],
    component_scores: Dict[str, Dict[str, float]],
    weights: Optional[Dict[str, float]] = None,
) -> List[RankedScenario]:
    """Deterministic, decomposed scenario ranking for REVIEW ordering only.

    component_scores: component name -> {scenario_id -> score in [0, 1]}.
    The returned score keeps its components attached so every ranking is
    explainable. Ties break on scenario_id, so the order is fully
    deterministic regardless of dict insertion order.

    IMPORTANT: this ordering must never be used to SELECT an evaluation set —
    "most interesting first" selection biases P(failure) estimates upward.
    Selection goes through stratified_sample(); ranking only orders items
    already selected (or a human review queue).
    """
    weights = weights or {name: 1.0 for name in component_scores}
    ranked = []
    for sid in scenario_ids:
        comps = {name: float(scores.get(sid, 0.0))
                 for name, scores in component_scores.items()}
        total_w = sum(weights.get(n, 0.0) for n in comps) or 1.0
        score = sum(comps[n] * weights.get(n, 0.0) for n in comps) / total_w
        ranked.append(RankedScenario(scenario_id=sid, score=round(score, 9),
                                     components=comps))
    ranked.sort(key=lambda r: (-r.score, r.scenario_id))
    return ranked