Store a quantile grid instead of a raw reservoir sample in the setup cache

NormalizerEntry.energy_reservoir_sample kept 100k raw energy values purely
to seed EnergyRouter centers via np.quantile at load time, which alone
accounted for most of the setup cache sidecar's ~2MB size (float32 values
round-tripped through Python floats serialize at full double precision).
Only a handful of quantile levels are ever read back, so collapse the
sample to a fixed 1001-point quantile grid at save time and interpolate
arbitrary levels from it at use time instead — about 100x smaller with
negligible (<0.001) error on the levels that matter. Bumps the cache
format version since old sidecars have no such grid to fall back on.

giant/data/setup_cache.py

@@ -31,7 +31,10 @@ from giant.data.transforms import Normalizer, sorted_membership
 # v2: event_id is now offset per-file (see loader.event_id_offset) to avoid
 # cross-file collisions, so a v1 sidecar's event_index/normalizers were
 # computed against collided ids and must not be reused.
-_CACHE_FORMAT_VERSION = 2
+# v3: NormalizerEntry.energy_reservoir_sample (100k raw values) replaced by
+# energy_quantiles (a fixed ENERGY_QUANTILE_LEVELS-point quantile grid) — a
+# v2 sidecar has no such grid to fall back on, so it must be recomputed.
+_CACHE_FORMAT_VERSION = 3
 
 _DIMS = {
     "COND_DIM": COND_DIM,
@@ -41,6 +44,29 @@ _DIMS = {
     "SEC_SLOT_DIM": SEC_SLOT_DIM,
 }
 
+# Resolution of the stored energy-quantile summary (see NormalizerEntry).
+# Only a handful of quantile *levels* (one per EnergyRouter expert) are ever
+# consumed (see pipeline.py), so a dense fixed grid of quantile values is
+# enough to reconstruct any level via interpolation (energy_quantile_at) —
+# at roughly 1/100th the storage of the raw 100k-value reservoir sample it
+# replaces, with negligible loss of resolution for that use.
+ENERGY_QUANTILE_LEVELS = 1001
+
+
+def energy_quantiles_from_sample(sample: np.ndarray) -> np.ndarray:
+    """Collapse a raw reservoir sample into the fixed grid stored on disk."""
+    if sample.size == 0:
+        return np.empty(0, dtype=np.float32)
+    levels = np.linspace(0.0, 1.0, ENERGY_QUANTILE_LEVELS)
+    return np.quantile(sample, levels).astype(np.float32)
+
+
+def energy_quantile_at(energy_quantiles: np.ndarray, levels: np.ndarray) -> np.ndarray:
+    """Interpolate quantile values at arbitrary probability `levels` from the
+    stored grid (e.g. `np.linspace(0, 1, n_experts)` for router centers)."""
+    grid_levels = np.linspace(0.0, 1.0, len(energy_quantiles))
+    return np.interp(levels, grid_levels, energy_quantiles).astype(np.float32)
+
 
 def sidecar_path(data: str | Path) -> Path:
     """The cache sidecar for `data`, always a sibling of `data` itself.
@@ -81,7 +107,10 @@ class NormalizerEntry:
     tgt_norm: Normalizer
     sec_phys_norm: Normalizer
     n_train_steps: int
-    energy_reservoir_sample: np.ndarray
+    energy_quantiles: np.ndarray
+    """Fixed ENERGY_QUANTILE_LEVELS-point quantile grid of the raw (pre-
+    normalization) pre-step energy column — see energy_quantiles_from_sample
+    / energy_quantile_at."""
 
     def to_json(self) -> dict:
         return {
@@ -89,8 +118,8 @@ class NormalizerEntry:
             "tgt_norm": self.tgt_norm.to_dict(),
             "sec_phys_norm": self.sec_phys_norm.to_dict(),
             "n_train_steps": self.n_train_steps,
-            "energy_reservoir_sample": np.asarray(
-                self.energy_reservoir_sample, dtype=np.float32
+            "energy_quantiles": np.asarray(
+                self.energy_quantiles, dtype=np.float32
             ).tolist(),
         }
 
@@ -101,9 +130,7 @@ class NormalizerEntry:
             tgt_norm=Normalizer.from_dict(d["tgt_norm"]),
             sec_phys_norm=Normalizer.from_dict(d["sec_phys_norm"]),
             n_train_steps=int(d["n_train_steps"]),
-            energy_reservoir_sample=np.array(
-                d["energy_reservoir_sample"], dtype=np.float32
-            ),
+            energy_quantiles=np.array(d["energy_quantiles"], dtype=np.float32),
         )