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 <noreply@anthropic.com>

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)

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]

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