Store a quantile grid instead of a raw reservoir sample in the setup cache
CI / Format (ruff format) (push) Successful in 26s
CI / Lint (ruff check) (push) Successful in 28s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 24s
CI / Tests (push) Successful in 58s
CI / Format (ruff format) (pull_request) Successful in 28s
CI / Lint (ruff check) (pull_request) Successful in 28s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 28s
CI / Tests (pull_request) Successful in 59s

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.
This commit is contained in:
2026-07-30 13:32:46 +02:00
parent de5db25e3f
commit d656cf3109
4 changed files with 87 additions and 26 deletions
+34 -1
View File
@@ -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 ──────────────────────────────────────────────