"""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