From d656cf3109da0c7a1533f0a024a5047891ff0e41 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 30 Jul 2026 13:32:46 +0200 Subject: [PATCH] Store a quantile grid instead of a raw reservoir sample in the setup cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 41 ++++++++++++++++++++++++++++++------- giant/pipeline.py | 35 ++++++++++++++++--------------- scripts/warm_setup_cache.py | 2 +- tests/test_setup_cache.py | 35 ++++++++++++++++++++++++++++++- 4 files changed, 87 insertions(+), 26 deletions(-) diff --git a/giant/data/setup_cache.py b/giant/data/setup_cache.py index f0a707f..229a99f 100644 --- a/giant/data/setup_cache.py +++ b/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), ) diff --git a/giant/pipeline.py b/giant/pipeline.py index c2de00e..d8c022d 100644 --- a/giant/pipeline.py +++ b/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']}" ) diff --git a/scripts/warm_setup_cache.py b/scripts/warm_setup_cache.py index 5885069..3eeebb3 100644 --- a/scripts/warm_setup_cache.py +++ b/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. """ diff --git a/tests/test_setup_cache.py b/tests/test_setup_cache.py index 04f5c00..5e1c7bd 100644 --- a/tests/test_setup_cache.py +++ b/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 ────────────────────────────────────────────── -- 2.39.5