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),
         )
 
 

giant/pipeline.py

@@ -149,7 +149,7 @@ def run_setup_stage(
         cond_norm = entry.cond_norm
         tgt_norm = entry.tgt_norm
         sec_phys_norm = entry.sec_phys_norm
-        energy_sample = entry.energy_reservoir_sample
+        energy_quantiles = entry.energy_quantiles
     else:
         echo("fitting normalizer (streaming) …")
         cond_acc = _WelfordAccumulator(COND_DIM)
@@ -158,12 +158,13 @@ def run_setup_stage(
         # EnergyRouter's default center spread (linspace over [-2, 2]) assumes
         # the z-normalized energy column is roughly uniform, which real energy
         # spectra rarely are — collect a reservoir sample here (reusing this
-        # same pass, not a second scan) so centers can instead be seeded from
-        # actual data quantiles below. Collected whenever the setup cache is
-        # being populated, not only when *this* run's router is
-        # energy-typed, so a later run enabling --router-type energy against
-        # this same (val_fraction, seed, conditioning) key never needs to
-        # rescan just to seed centers.
+        # same pass, not a second scan), then collapse it to a fixed quantile
+        # grid (setup_cache.energy_quantiles_from_sample) so centers can
+        # instead be seeded from actual data quantiles below. Collected
+        # whenever the setup cache is being populated, not only when *this*
+        # run's router is energy-typed, so a later run enabling
+        # --router-type energy against this same (val_fraction, seed,
+        # conditioning) key never needs to rescan just to seed centers.
         collect_energy_sample = energy_router_active or cache is not None
         energy_sampler = (
             _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
@@ -194,24 +195,24 @@ def run_setup_stage(
         cond_norm = cond_acc.to_normalizer()
         tgt_norm = tgt_acc.to_normalizer()
         sec_phys_norm = sec_phys_acc.to_normalizer()
-        energy_sample = (
-            energy_sampler.sample
+        energy_quantiles = (
+            setup_cache.energy_quantiles_from_sample(energy_sampler.sample)
             if energy_sampler is not None
             else np.empty(0, dtype=np.float32)
         )
         if cache is not None:
             cache.normalizers[norm_key] = setup_cache.NormalizerEntry(
-                cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_sample
+                cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_quantiles
             )
 
-    if energy_router_active and energy_sample.size > 0:
+    if energy_router_active and energy_quantiles.size > 0:
         assert cond_norm.mean is not None and cond_norm.std is not None
-        normalized_sample = (
-            energy_sample - cond_norm.mean[energy_idx]
-        ) / cond_norm.std[energy_idx]
-        quantiles = np.linspace(0.0, 1.0, router_cfg["n_experts"])
-        centers_init = np.quantile(normalized_sample, quantiles).astype(np.float32)
-        router_cfg["centers_init"] = centers_init.tolist()
+        levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
+        raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
+        centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[
+            energy_idx
+        ]
+        router_cfg["centers_init"] = centers_init.astype(np.float32).tolist()
         echo(
             f"  seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
         )

scripts/warm_setup_cache.py

@@ -30,7 +30,7 @@ def run_warm_setup_cache(
     `giant train` invocation will use so it hits this warmed entry.
     `router_enabled`/`router_type`/`n_experts` only matter for
     `router_type == "process"` (warms that `n_experts`'s process map); the
-    energy-router reservoir sample is always collected regardless, so a
+    energy-router quantile summary is always collected regardless, so a
     later `--router-type energy` run never needs to rescan just to seed
     centers.
     """

tests/test_setup_cache.py

@@ -98,7 +98,7 @@ def test_save_load_round_trip(tmp_path):
     assert entry.cond_norm.mean is not None
     np.testing.assert_allclose(entry.cond_norm.mean, np.zeros(3, dtype=np.float32))
     assert entry.n_train_steps == 100
-    np.testing.assert_allclose(entry.energy_reservoir_sample, [1.0, 2.0, 3.0])
+    np.testing.assert_allclose(entry.energy_quantiles, [1.0, 2.0, 3.0])
 
 
 def test_load_missing_sidecar_returns_none(tmp_path):
@@ -214,6 +214,39 @@ def test_save_merges_non_colliding_normalizer_keys(tmp_path):
     assert loaded.normalizers["k2"].n_train_steps == 2
 
 
+# ── energy_quantiles_from_sample / energy_quantile_at ───────────────────
+
+
+def test_energy_quantiles_from_sample_empty():
+    result = setup_cache.energy_quantiles_from_sample(np.empty(0, dtype=np.float32))
+    assert result.size == 0
+
+
+def test_energy_quantiles_from_sample_has_fixed_grid_size():
+    sample = np.random.default_rng(0).normal(size=5000).astype(np.float32)
+    result = setup_cache.energy_quantiles_from_sample(sample)
+    assert result.shape == (setup_cache.ENERGY_QUANTILE_LEVELS,)
+    assert result[0] == pytest.approx(sample.min(), abs=1e-3)
+    assert result[-1] == pytest.approx(sample.max(), abs=1e-3)
+
+
+def test_energy_quantile_at_matches_direct_quantile_on_stored_grid():
+    sample = np.random.default_rng(1).exponential(size=20_000).astype(np.float32)
+    grid = setup_cache.energy_quantiles_from_sample(sample)
+
+    levels = np.linspace(0.0, 1.0, 5)
+    got = setup_cache.energy_quantile_at(grid, levels)
+    expected = np.quantile(sample, levels)
+
+    np.testing.assert_allclose(got, expected, rtol=0.05)
+
+
+def test_energy_quantile_at_median_of_two_points():
+    grid = np.array([0.0, 10.0], dtype=np.float32)
+    result = setup_cache.energy_quantile_at(grid, np.array([0.0, 0.5, 1.0]))
+    np.testing.assert_allclose(result, [0.0, 5.0, 10.0])
+
+
 # ── n_train_steps_for_split ──────────────────────────────────────────────