Apply ruff format after merging phase2-secondary-prediction

The merged proc_idx/proc_map plumbing wasn't run through ruff format
before merging; reflow only, no logic changes.
This commit is contained in:
2026-07-15 10:00:36 +02:00
parent a5683517be
commit 05d5dee606
7 changed files with 67 additions and 26 deletions
+40 -12
View File
@@ -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,
)
+1 -3
View File
@@ -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, ...)
+10 -3
View File
@@ -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,
)
)
+6 -2
View File
@@ -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.")
+1 -3
View File
@@ -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)
+3 -1
View File
@@ -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]]
+6 -2
View File
@@ -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}