From 24445b7427fe9b5d0becced3175d0bc32fefa513 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 10 Aug 2026 11:38:04 +0200 Subject: [PATCH] Add coverage for router-center seeding, geometry batch reader, material topN cache, and setup-cache corruption paths Closes the highest-value coverage gaps found via pytest-cov: pipeline.py's EnergyRouter quantile-seeding (the roadmap's flagged fix for the failed MoE rollout benchmark) had zero coverage, geometry.py's real parquet-batch reader was always mocked, the material top-N-map cache-hit branch was untested (only pdg's was), and setup_cache.py was missing malformed-cache-body and unknown-axis error paths. Co-Authored-By: Claude Sonnet 5 --- tests/test_geometry.py | 71 ++++++++++++++++++++++++++++++++ tests/test_pipeline.py | 86 ++++++++++++++++++++++++++++++++++++++- tests/test_setup_cache.py | 32 +++++++++++++++ 3 files changed, 188 insertions(+), 1 deletion(-) diff --git a/tests/test_geometry.py b/tests/test_geometry.py index d967315..9e9aace 100644 --- a/tests/test_geometry.py +++ b/tests/test_geometry.py @@ -4,6 +4,7 @@ from pathlib import Path from unittest.mock import patch import numpy as np +import pandas as pd import pytest from giant import geometry as g @@ -12,6 +13,76 @@ from giant import geometry as g pytest.importorskip("sklearn") +def _steps_frame(n=5, with_post=True): + rng = np.random.default_rng(0) + data = { + "pre_x": rng.uniform(-10, 10, n), + "pre_y": rng.uniform(-10, 10, n), + "pre_z": rng.uniform(-10, 10, n), + "material": ["G4_AIR"] * n, + "layer_id": np.arange(n, dtype=np.int64), + } + if with_post: + data["post_x"] = rng.uniform(-10, 10, n) + data["post_y"] = rng.uniform(-10, 10, n) + data["post_z"] = rng.uniform(-10, 10, n) + return pd.DataFrame(data) + + +def test_iter_point_batches_missing_columns_raises(tmp_path): + path = tmp_path / "steps.parquet" + pd.DataFrame({"pre_x": [0.0]}).to_parquet(path) + with pytest.raises(ValueError, match="missing columns"): + next(g._iter_point_batches(path)) + + +def test_iter_point_batches_without_post_columns_yields_pre_only(tmp_path): + path = tmp_path / "steps.parquet" + df = _steps_frame(n=5, with_post=False) + df.to_parquet(path) + + (pos, mat, lay) = next(g._iter_point_batches(path)) + + assert pos.shape == (5, 3) + np.testing.assert_allclose( + pos, df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32) + ) + assert list(mat) == ["G4_AIR"] * 5 + np.testing.assert_array_equal(lay, np.arange(5)) + + +def test_iter_point_batches_with_post_columns_doubles_and_concatenates_points( + tmp_path, +): + path = tmp_path / "steps.parquet" + df = _steps_frame(n=5, with_post=True) + df.to_parquet(path) + + (pos, mat, lay) = next(g._iter_point_batches(path)) + + # Every step contributes both its pre_pos and post_pos, sharing the + # step's material/layer_id label — so batches double in length. + assert pos.shape == (10, 3) + np.testing.assert_allclose( + pos[:5], df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32) + ) + np.testing.assert_allclose( + pos[5:], df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32) + ) + assert list(mat) == ["G4_AIR"] * 10 + np.testing.assert_array_equal(lay, np.concatenate([np.arange(5), np.arange(5)])) + + +def test_iter_point_batches_respects_batch_size(tmp_path): + path = tmp_path / "steps.parquet" + df = _steps_frame(n=10, with_post=False) + df.to_parquet(path, row_group_size=10) + + batches = list(g._iter_point_batches(path, batch_size=4)) + + assert [len(pos) for pos, _, _ in batches] == [4, 4, 2] + + def _box_batch(n, rng): """A labelled point cloud: inside a 100mm box -> PbWO4/0, else AIR/-1.""" pos = rng.uniform(-200, 200, (n, 3)).astype(np.float32) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index fae62c5..c0848c9 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -6,9 +6,10 @@ import pytest import torch from giant import config as gconfig +from giant.constants import COND_DIM from giant.data import setup_cache from giant.data.transforms import Normalizer -from giant.pipeline import run_train_job +from giant.pipeline import _seed_energy_router, run_train_job def _unit(v): @@ -175,6 +176,29 @@ def test_run_train_job_builds_caches_and_persists_pdg_topn_map(tmp_path, data): assert any("pdg top-N map: cache hit" in m for m in echo2) +def test_run_train_job_builds_caches_and_persists_material_topn_map(tmp_path, data): + """conditioning.material.type="onehot" is an independent axis from the + pdg one above, with its own build/cache-hit branch in run_setup_stage — + exercise both here the same way the pdg test above does.""" + cfg = _tiny_cfg() + cfg["conditioning"]["material"]["type"] = "onehot" + echo1 = _run(data, tmp_path / "out1", cfg=cfg) + assert any("building material top-N map" in m for m in echo1) + + loaded = setup_cache.load(data, [data]) + assert loaded is not None + key = setup_cache.topn_key("material", 4) # conditioning.material.emb_dim = 4 + assert key in loaded.topn_maps + assert set(loaded.topn_maps[key].class_map.keys()) >= {"G4_AIR", "G4_Fe"} + + ckpt = torch.load(tmp_path / "out1" / "last.pt", weights_only=False) + assert "mat_topn_map" in ckpt + assert set(ckpt["mat_topn_map"]["class_map"].keys()) >= {"G4_AIR", "G4_Fe"} + + echo2 = _run(data, tmp_path / "out2", cfg=cfg) + assert any("material top-N map: cache hit" in m for m in echo2) + + def test_run_train_job_no_topn_map_for_physical_target(tmp_path, data): cfg = _tiny_cfg() cfg["stage2_model"]["particle_type"] = {"target": "physical", "lambda": 1.0} @@ -326,3 +350,63 @@ def test_run_train_job_matches_uncached_output(tmp_path, data): ) assert uncached["pdg_map"] == cached["pdg_map"] assert uncached["mat_map"] == cached["mat_map"] + + +def _fitted_cond_norm(seed=0): + rng = np.random.default_rng(seed) + return Normalizer().fit(rng.normal(size=(64, COND_DIM)).astype(np.float32)) + + +@pytest.mark.parametrize( + "router_cfg", + [ + {"enabled": False, "type": "energy", "n_experts": 4}, + {"enabled": True, "type": "pdg", "n_experts": 4}, + ], +) +def test_seed_energy_router_noop_when_not_an_enabled_energy_router(router_cfg): + cond_norm = _fitted_cond_norm() + echoed = [] + _seed_energy_router(router_cfg, cond_norm, np.array([1.0, 2.0]), 3, echoed.append) + assert "centers_init" not in router_cfg + assert echoed == [] + + +def test_seed_energy_router_falls_back_to_default_and_warns_when_no_samples(): + router_cfg = {"enabled": True, "type": "energy", "n_experts": 4} + cond_norm = _fitted_cond_norm() + echoed = [] + _seed_energy_router( + router_cfg, cond_norm, np.empty(0), energy_idx=3, echo=echoed.append + ) + assert "centers_init" not in router_cfg + assert len(echoed) == 1 + assert "falls back to default centers" in echoed[0] + + +def test_seed_energy_router_seeds_centers_from_data_quantiles(): + router_cfg = {"enabled": True, "type": "energy", "n_experts": 4} + cond_norm = _fitted_cond_norm() + energy_idx = 3 + # A grid of "raw" quantile values as setup_cache.energy_quantiles_from_sample + # would produce them: monotonically increasing, in the same (log-energy) + # units as the conditioning column being normalized against. + energy_quantiles = np.linspace(1.0, 10.0, 33).astype(np.float32) + echoed = [] + _seed_energy_router( + router_cfg, cond_norm, energy_quantiles, energy_idx, echoed.append + ) + + assert "centers_init" in router_cfg + centers = np.asarray(router_cfg["centers_init"], dtype=np.float32) + assert centers.shape == (router_cfg["n_experts"],) + + levels = np.linspace(0.0, 1.0, router_cfg["n_experts"]) + raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels) + expected = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[energy_idx] + np.testing.assert_allclose(centers, expected, rtol=1e-5) + # Quantile levels are increasing, and the normalizer's std is positive, so + # the seeded centers must preserve that order rather than e.g. reversing it. + assert np.all(np.diff(centers) > 0) + assert len(echoed) == 1 + assert "seeded EnergyRouter centers" in echoed[0] diff --git a/tests/test_setup_cache.py b/tests/test_setup_cache.py index 9224b64..70e8e40 100644 --- a/tests/test_setup_cache.py +++ b/tests/test_setup_cache.py @@ -128,6 +128,11 @@ def test_save_load_round_trip_topn_maps(tmp_path): assert mat_m.class_map == {"G4_AIR": 0, "PbWO4": 1} +def test_topn_key_unknown_axis_raises(): + with pytest.raises(ValueError, match="unknown top-N map axis"): + setup_cache.topn_key("process", 4) + + def test_load_missing_sidecar_returns_none(tmp_path): data = _touch_parquet(tmp_path / "shard.parquet") assert setup_cache.load(data, [data]) is None @@ -176,6 +181,25 @@ def test_load_invalidates_on_file_content_change(tmp_path): assert setup_cache.load(data, files) is None +def test_load_returns_none_on_malformed_cache_body(tmp_path): + """format_version/dims/fingerprint all check out, but the cache body + itself doesn't match SetupCache.from_json's expected shape (e.g. hand- + edited or written by a version that changed a nested key) — a clean + miss, not a crash.""" + data = _touch_parquet(tmp_path / "shard.parquet") + files = [data] + setup_cache.save(data, files, SetupCache.empty(files)) + + path = setup_cache.sidecar_path(data) + raw = json.loads(path.read_text()) + raw["vocab"] = {"pdg_map": {"11": 0}} # missing required "mat_map" key + path.write_text(json.dumps(raw)) + + echoed = [] + assert setup_cache.load(data, files, echo=echoed.append) is None + assert any("malformed" in m for m in echoed) + + def test_load_soft_warns_on_git_hash_mismatch_but_still_hits(tmp_path, capsys): data = _touch_parquet(tmp_path / "shard.parquet") files = [data] @@ -352,3 +376,11 @@ def test_compute_event_index_from_files_single_file_unaffected(tmp_path): np.testing.assert_array_equal(unique_ids, [5, 7]) np.testing.assert_array_equal(counts, [2, 1]) + + +def test_compute_event_index_from_files_empty_file_list(): + unique_ids, counts = setup_cache.compute_event_index_from_files([]) + assert unique_ids.size == 0 + assert counts.size == 0 + assert unique_ids.dtype == np.int64 + assert counts.dtype == np.int64