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:
+40
-12
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
@@ -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,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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]]
|
||||
|
||||
@@ -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}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user