diff --git a/giant/data/dataset.py b/giant/data/dataset.py index afdb051..f319d93 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -100,8 +100,15 @@ class StreamingStepsDataset(IterableDataset): chunk = {k: v[mask] for k, v in chunk.items()} ( - cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, - proc_idx, _, _, + cond_cont, + cond_cat, + target_s1, + n_sec, + sec_cont, + sec_pdg_idx, + proc_idx, + _, + _, ) = build_features( chunk, self.pdg_map, @@ -122,18 +129,34 @@ class StreamingStepsDataset(IterableDataset): if buf_n >= self.shuffle_buffer: ( - buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, - buf_proc, buf_n, - ) = ( - yield from self._flush( - buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, - buf_proc, final=False, - ) + buf_cont, + buf_cat, + buf_tgt, + buf_nsec, + buf_sec, + buf_spdg, + buf_proc, + buf_n, + ) = yield from self._flush( + buf_cont, + buf_cat, + buf_tgt, + buf_nsec, + buf_sec, + buf_spdg, + buf_proc, + final=False, ) if buf_n > 0: yield from self._flush( - buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, buf_proc, + buf_cont, + buf_cat, + buf_tgt, + buf_nsec, + buf_sec, + buf_spdg, + buf_proc, final=True, ) @@ -180,7 +203,12 @@ class StreamingStepsDataset(IterableDataset): return [], [], [], [], [], [], [], 0 rem = n_full * bs return ( - [cont[rem:]], [cat[rem:]], [tgt[rem:]], - [nsec[rem:]], [sec[rem:]], [spdg[rem:]], [proc[rem:]], + [cont[rem:]], + [cat[rem:]], + [tgt[rem:]], + [nsec[rem:]], + [sec[rem:]], + [spdg[rem:]], + [proc[rem:]], n - rem, ) diff --git a/giant/data/loader.py b/giant/data/loader.py index 86dd989..4ddbf7f 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -202,9 +202,7 @@ def build_index_maps_from_files( ) -def build_process_map_from_files( - files: list[Path], n_experts: int -) -> dict[str, int]: +def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str, int]: """Scan the `process` column and build a frequency-capped process->index map. Physics processes have a long tail (rare nuclear captures, decays, ...) diff --git a/giant/pipeline.py b/giant/pipeline.py index 2a0ba96..f160231 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -59,8 +59,12 @@ def run_train_job( proc_map: dict[str, int] | None = None if router_cfg.get("enabled") and router_cfg.get("type") == "process": echo("building process vocabulary …") - proc_map = build_process_map_from_files(files, n_experts=router_cfg["n_experts"]) - echo(f" {len(proc_map)} process labels mapped to {router_cfg['n_experts']} experts") + proc_map = build_process_map_from_files( + files, n_experts=router_cfg["n_experts"] + ) + echo( + f" {len(proc_map)} process labels mapped to {router_cfg['n_experts']} experts" + ) echo("fitting normalizer (streaming) …") cond_acc = _WelfordAccumulator(COND_DIM) @@ -73,7 +77,10 @@ def run_train_job( chunk_tr = {k: v[mask] for k, v in chunk.items()} cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _proc, _, _ = ( build_features( - chunk_tr, pdg_map, mat_map, proc_map=proc_map, + chunk_tr, + pdg_map, + mat_map, + proc_map=proc_map, require_secondaries=True, ) ) diff --git a/scripts/steps_to_parquet_parallel.py b/scripts/steps_to_parquet_parallel.py index dc1c5fe..b2ee41e 100644 --- a/scripts/steps_to_parquet_parallel.py +++ b/scripts/steps_to_parquet_parallel.py @@ -217,8 +217,12 @@ def run_parallel_job( raise SystemExit(1) total_orphaned = sum( - int(m.group(1)) for _, _, stdout, _ in results for m in _ORPHAN_RE.finditer(stdout) + int(m.group(1)) + for _, _, stdout, _ in results + for m in _ORPHAN_RE.finditer(stdout) ) if total_orphaned: - print(f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s).") + print( + f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s)." + ) print(f"\nAll {len(results)} conversion(s) completed.") diff --git a/tests/test_loader.py b/tests/test_loader.py index 703d9d3..af609dc 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -65,9 +65,7 @@ def test_build_process_map_from_files_keeps_most_frequent(tmp_path): """process counts: eIoni=5, phot=3, compt=2, Rayl=1 — with n_experts=3, only the top 2 (eIoni, phot) get their own index; compt/Rayl share the "other" (last) index.""" - process = ( - ["eIoni"] * 5 + ["phot"] * 3 + ["compt"] * 2 + ["Rayl"] * 1 - ) + process = ["eIoni"] * 5 + ["phot"] * 3 + ["compt"] * 2 + ["Rayl"] * 1 path = tmp_path / "shard-000.parquet" pd.DataFrame({"process": process}).to_parquet(path) diff --git a/tests/test_steps_to_parquet.py b/tests/test_steps_to_parquet.py index 9810d6e..8e839c1 100644 --- a/tests/test_steps_to_parquet.py +++ b/tests/test_steps_to_parquet.py @@ -69,7 +69,9 @@ def test_orphaned_child_track_is_dropped_not_nulled(): } ) out, n_orphaned = steps_to_parquet._add_secondary_attributes(df) - row = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0)) + row = out.filter( + (pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0) + ) assert n_orphaned == 1 assert row["child_track_ids"].to_list() == [[2]] diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 86ab26d..c95a36b 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -272,7 +272,9 @@ def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict: def test_build_features_proc_idx_zero_without_proc_map(): - data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object)) + data = _minimal_step_data( + 3, process=np.array(["compt", "phot", "eIoni"], dtype=object) + ) pdg_map, mat_map = {11: 0}, {"PbWO4": 0} *_, proc_idx, _, _ = build_features(data, pdg_map, mat_map) @@ -281,7 +283,9 @@ def test_build_features_proc_idx_zero_without_proc_map(): def test_build_features_proc_idx_looks_up_proc_map(): - data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object)) + data = _minimal_step_data( + 3, process=np.array(["compt", "phot", "eIoni"], dtype=object) + ) pdg_map, mat_map = {11: 0}, {"PbWO4": 0} proc_map = {"compt": 0, "phot": 1, "eIoni": 2}