sensorflow/megaeval/population.py

"""Immutable synthetic populations, generated vectorized and stored partitioned.

Storage layout (single-node stand-in for an Iceberg/parquet lakehouse — no parquet
engine is installed in this venv, so partitions are compressed npz column files):

    runs/megaeval/populations/{pop_id}/
        meta.json          dataset card: dims vocab, counts (exact), train mix, lineage
        part-0000.npz ...  column arrays partitioned by container_id % NUM_PARTITIONS

Every object row carries:
  - 9 low-cardinality dimension codes (the cube dimensions)
  - a model-independent `difficulty` in [0,1] (occlusion/distance/lighting/weather/size)
  - `safety_critical` flag (VRUs in high-risk scenarios)

The "training mix" is a deliberately shifted joint distribution over
(class, weather, lighting) used by shift detection: training saw less
night/rain/fog than this evaluation population.
"""

from __future__ import annotations

import json
import os
import threading
import uuid
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple

import numpy as np

MEGA_ROOT = os.path.join("runs", "megaeval")
NUM_PARTITIONS = 32

DIMENSIONS: Dict[str, List[str]] = {
    "class": ["vehicle", "pedestrian", "cyclist", "motorcycle", "truck", "bus"],
    "weather": ["clear", "rain", "fog", "snow"],
    "lighting": ["day", "night", "dusk"],
    "road_type": ["highway", "urban", "rural", "intersection"],
    "scenario": ["nominal", "near_miss", "occlusion_heavy", "night_glare",
                 "extreme_ttc", "sensor_degraded"],
    "sensor": ["lidar", "camera", "fused"],
    "distance_band": ["0-20m", "20-50m", "50-100m", "100m+"],
    "speed_band": ["stationary", "slow", "medium", "fast"],
    "occlusion": ["none", "partial", "heavy"],
}
DIM_NAMES = list(DIMENSIONS.keys())
CONTAINER_DIMS = ["weather", "lighting", "road_type", "scenario"]
OBJECT_DIMS = ["class", "sensor", "distance_band", "speed_band", "occlusion"]


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def pop_dir(pop_id: str) -> str:
    return os.path.join(MEGA_ROOT, "populations", pop_id)


def set_mega_root(path: str) -> None:
    """Test hook: relocate all megaeval storage (and clear partition caches)."""
    global MEGA_ROOT
    MEGA_ROOT = str(path)
    with _CACHE_LOCK:
        _PART_CACHE.clear()


# ------------------------------------------------------------------ generation


def _sample_container_dims(rng: np.random.Generator, n: int) -> Dict[str, np.ndarray]:
    """Correlated container-level context (weather/lighting/road/scenario)."""
    weather = rng.choice(4, size=n, p=[0.62, 0.20, 0.11, 0.07]).astype(np.int8)
    lighting = rng.choice(3, size=n, p=[0.60, 0.28, 0.12]).astype(np.int8)
    road = rng.choice(4, size=n, p=[0.30, 0.38, 0.14, 0.18]).astype(np.int8)
    scenario = rng.choice(6, size=n, p=[0.66, 0.09, 0.09, 0.06, 0.05, 0.05]).astype(np.int8)
    # correlations: night_glare scenario forces night; fog/snow raise sensor_degraded odds
    scenario_names = DIMENSIONS["scenario"]
    lighting[scenario == scenario_names.index("night_glare")] = 1  # night
    bad_wx = (weather >= 2)
    flip = bad_wx & (rng.random(n) < 0.18)
    scenario[flip] = scenario_names.index("sensor_degraded")
    return {"weather": weather, "lighting": lighting, "road_type": road, "scenario": scenario}


_CLASS_MIX_BY_ROAD = np.array([
    # vehicle, pedestrian, cyclist, motorcycle, truck, bus
    [0.62, 0.02, 0.01, 0.05, 0.24, 0.06],   # highway
    [0.44, 0.26, 0.12, 0.06, 0.07, 0.05],   # urban
    [0.55, 0.08, 0.06, 0.08, 0.18, 0.05],   # rural
    [0.40, 0.30, 0.14, 0.05, 0.06, 0.05],   # intersection
])

_SPEED_MIX_BY_ROAD = np.array([
    [0.02, 0.08, 0.35, 0.55],   # highway
    [0.22, 0.44, 0.28, 0.06],   # urban
    [0.10, 0.30, 0.40, 0.20],   # rural
    [0.30, 0.45, 0.20, 0.05],   # intersection
])


def _object_difficulty(codes: Dict[str, np.ndarray], rng: np.random.Generator) -> np.ndarray:
    """Model-independent difficulty in [0,1] — the physical hardness of each object."""
    n = codes["class"].shape[0]
    d = np.full(n, 0.12, dtype=np.float32)
    d += np.array([0.00, 0.14, 0.12, 0.10, 0.02, 0.02], dtype=np.float32)[codes["class"]]
    d += np.array([0.00, 0.08, 0.12, 0.10], dtype=np.float32)[codes["weather"]]
    d += np.array([0.00, 0.12, 0.05], dtype=np.float32)[codes["lighting"]]
    d += np.array([0.00, 0.10, 0.24], dtype=np.float32)[codes["occlusion"]]
    d += np.array([0.00, 0.05, 0.14, 0.26], dtype=np.float32)[codes["distance_band"]]
    d += np.array([0.00, 0.04, 0.08, 0.06, 0.07, 0.12], dtype=np.float32)[codes["scenario"]]
    d += rng.normal(0.0, 0.05, size=n).astype(np.float32)
    return np.clip(d, 0.01, 0.99)


def generate_population(
    name: str = "mega-perception",
    num_objects: int = 320_000,
    seed: int = 42,
    avg_objects_per_container: float = 12.0,
) -> Dict:
    """Generate + persist a partitioned population. Vectorized; ~seconds for 320k."""
    rng = np.random.default_rng(seed)
    num_containers = max(64, int(num_objects / avg_objects_per_container))

    cdims = _sample_container_dims(rng, num_containers)
    per_container = np.clip(rng.poisson(avg_objects_per_container, size=num_containers),
                            3, 48).astype(np.int64)
    total = int(per_container.sum())
    container_id = np.repeat(np.arange(num_containers, dtype=np.int64), per_container)

    codes: Dict[str, np.ndarray] = {k: cdims[k][container_id] for k in CONTAINER_DIMS}

    road_obj = codes["road_type"].astype(np.int64)
    u = rng.random(total)
    cls_cdf = np.cumsum(_CLASS_MIX_BY_ROAD, axis=1)[road_obj]
    codes["class"] = (u[:, None] > cls_cdf).sum(axis=1).astype(np.int8)
    u2 = rng.random(total)
    spd_cdf = np.cumsum(_SPEED_MIX_BY_ROAD, axis=1)[road_obj]
    codes["speed_band"] = (u2[:, None] > spd_cdf).sum(axis=1).astype(np.int8)
    codes["sensor"] = rng.choice(3, size=total, p=[0.38, 0.22, 0.40]).astype(np.int8)
    codes["distance_band"] = rng.choice(4, size=total, p=[0.24, 0.40, 0.25, 0.11]).astype(np.int8)

    occ_p = np.tile(np.array([0.62, 0.28, 0.10]), (total, 1))
    heavy_scene = codes["scenario"] == DIMENSIONS["scenario"].index("occlusion_heavy")
    occ_p[heavy_scene] = [0.20, 0.42, 0.38]
    occ_cdf = np.cumsum(occ_p, axis=1)
    codes["occlusion"] = (rng.random(total)[:, None] > occ_cdf).sum(axis=1).astype(np.int8)

    difficulty = _object_difficulty(codes, rng)
    vru = np.isin(codes["class"], [1, 2, 3])
    risky_scene = np.isin(codes["scenario"], [1, 4])  # near_miss, extreme_ttc
    near = codes["distance_band"] <= 1
    safety_critical = (vru & (risky_scene | (near & (rng.random(total) < 0.15)))).astype(np.uint8)

    object_id = np.arange(total, dtype=np.int64)
    pop_id = f"pop-{uuid.uuid4().hex[:10]}"
    d = pop_dir(pop_id)
    os.makedirs(d, exist_ok=True)

    part_of = (container_id % NUM_PARTITIONS).astype(np.int64)
    order = np.argsort(part_of, kind="stable")
    part_sorted = part_of[order]
    bounds = np.searchsorted(part_sorted, np.arange(NUM_PARTITIONS + 1))
    partition_meta = []
    for p in range(NUM_PARTITIONS):
        sl = order[bounds[p]:bounds[p + 1]]
        cols = {k: codes[k][sl] for k in DIM_NAMES}
        cols["object_id"] = object_id[sl]
        cols["container_id"] = container_id[sl]
        cols["difficulty"] = difficulty[sl]
        cols["safety_critical"] = safety_critical[sl]
        np.savez_compressed(os.path.join(d, f"part-{p:04d}.npz"), **cols)
        partition_meta.append({"partition": p, "rows": int(sl.size)})

    # exact per-dimension counts (dataset card)
    dim_counts = {
        k: {DIMENSIONS[k][i]: int(c) for i, c in enumerate(np.bincount(codes[k], minlength=len(DIMENSIONS[k])))}
        for k in DIM_NAMES
    }

    train_mix = _training_mix(codes, rng)

    meta = {
        "population_id": pop_id,
        "name": name,
        "created_at": _now(),
        "seed": seed,
        "num_objects": total,
        "num_containers": num_containers,
        "num_partitions": NUM_PARTITIONS,
        "avg_objects_per_container": round(total / num_containers, 2),
        "dimensions": DIMENSIONS,
        "dim_counts": dim_counts,
        "partitions": partition_meta,
        "safety_critical_count": int(safety_critical.sum()),
        "train_mix": train_mix,
    }
    with open(os.path.join(d, "meta.json"), "w") as f:
        json.dump(meta, f)
    return meta


def _training_mix(codes: Dict[str, np.ndarray], rng: np.random.Generator) -> List[Dict]:
    """Joint (class, weather, lighting) proportions the *training* set had.

    Derived from the eval mix with deliberate, deterministic under-representation of
    hard conditions (night, rain/fog/snow, VRU at night) so shift detection has
    honest signal: shift = eval share vs train share of the same cohort.
    """
    n = codes["class"].shape[0]
    key = (codes["class"].astype(np.int64) * 100
           + codes["weather"].astype(np.int64) * 10
           + codes["lighting"].astype(np.int64))
    uniq, counts = np.unique(key, return_counts=True)
    eval_share = counts / n

    cls, wx, lt = uniq // 100, (uniq // 10) % 10, uniq % 10
    factor = np.ones(uniq.shape)
    factor[lt == 1] *= 0.55                       # trained on less night
    factor[wx >= 1] *= 0.70                       # less rain/fog/snow
    factor[(np.isin(cls, [1, 2])) & (lt == 1)] *= 0.75  # far fewer night VRUs
    train_share = eval_share * factor
    train_share = train_share / train_share.sum()

    out = []
    for i in range(uniq.size):
        out.append({
            "class": DIMENSIONS["class"][int(cls[i])],
            "weather": DIMENSIONS["weather"][int(wx[i])],
            "lighting": DIMENSIONS["lighting"][int(lt[i])],
            "train_share": float(train_share[i]),
            "eval_share": float(eval_share[i]),
            "eval_count": int(counts[i]),
        })
    return out


# ------------------------------------------------------------------ loading


_PART_CACHE: Dict[Tuple[str, int], Dict[str, np.ndarray]] = {}
_CACHE_LOCK = threading.Lock()


def load_meta(pop_id: str) -> Optional[Dict]:
    path = os.path.join(pop_dir(pop_id), "meta.json")
    if not os.path.exists(path):
        return None
    with open(path) as f:
        return json.load(f)


def list_populations() -> List[Dict]:
    base = os.path.join(MEGA_ROOT, "populations")
    if not os.path.isdir(base):
        return []
    metas = []
    for pid in sorted(os.listdir(base)):
        m = load_meta(pid)
        if m:
            metas.append(m)
    return metas


def load_partition(pop_id: str, partition: int) -> Dict[str, np.ndarray]:
    key = (pop_id, partition)
    with _CACHE_LOCK:
        if key in _PART_CACHE:
            return _PART_CACHE[key]
    path = os.path.join(pop_dir(pop_id), f"part-{partition:04d}.npz")
    with np.load(path) as z:
        cols = {k: z[k] for k in z.files}
    with _CACHE_LOCK:
        if len(_PART_CACHE) > 96:  # bounded partition cache
            _PART_CACHE.clear()
        _PART_CACHE[key] = cols
    return cols


def partition_of_container(container_id: int) -> int:
    return int(container_id) % NUM_PARTITIONS