Add tooling for a versioned geant_steps dataset layout
Introduces raw/<kind>/<gen>/<detector>/shard-NNN.root and processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet as the dataset convention, plus scripts to operate on it: migrate_geant_steps.py for the one-time move into this layout, bump_dataset_version.py to cut new gen/schema versions with a logged reason, steps_to_parquet_parallel.py to convert ROOT shards to parquet in parallel and place them correctly, and create_root_files.py to generate new ROOT shards via a minicalosim executable. The loader gains .manifest file support so pools/ (train/dev/ holdout shard lists) can be passed straight to `giant train`. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,95 @@
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
|
||||
# scripts/ is not an installed package — load the module straight from its path.
|
||||
_SPEC = importlib.util.spec_from_file_location(
|
||||
"bump_dataset_version",
|
||||
Path(__file__).resolve().parents[1] / "scripts" / "bump_dataset_version.py",
|
||||
)
|
||||
bump_dataset_version = importlib.util.module_from_spec(_SPEC)
|
||||
_SPEC.loader.exec_module(bump_dataset_version)
|
||||
|
||||
plan_bump_gen = bump_dataset_version.plan_bump_gen
|
||||
plan_bump_schema = bump_dataset_version.plan_bump_schema
|
||||
apply_bump = bump_dataset_version.apply_bump
|
||||
|
||||
|
||||
def test_bump_gen_starts_at_gen1_when_none_exist(tmp_path):
|
||||
dirs, log_line = plan_bump_gen(tmp_path, "steps", "first generation", None, "2026-01-01")
|
||||
assert dirs == [
|
||||
tmp_path / "raw" / "steps" / "gen1",
|
||||
tmp_path / "processed" / "steps" / "gen1" / "schema1",
|
||||
]
|
||||
assert "`gen1`" in log_line
|
||||
assert "first generation" in log_line
|
||||
|
||||
|
||||
def test_bump_gen_increments_past_existing(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
(tmp_path / "raw" / "steps" / "gen2").mkdir(parents=True)
|
||||
dirs, _ = plan_bump_gen(tmp_path, "steps", "next gen", None, "2026-01-01")
|
||||
assert dirs[0] == tmp_path / "raw" / "steps" / "gen3"
|
||||
|
||||
|
||||
def test_bump_gen_checks_both_raw_and_processed_trees(tmp_path):
|
||||
# processed/ is ahead of raw/ — next gen must still be past the max of both.
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
(tmp_path / "processed" / "steps" / "gen4" / "schema1").mkdir(parents=True)
|
||||
dirs, _ = plan_bump_gen(tmp_path, "steps", "next gen", None, "2026-01-01")
|
||||
assert dirs[0] == tmp_path / "raw" / "steps" / "gen5"
|
||||
|
||||
|
||||
def test_bump_gen_kinds_are_independent(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen5").mkdir(parents=True)
|
||||
dirs, _ = plan_bump_gen(tmp_path, "hits", "first hits gen", None, "2026-01-01")
|
||||
assert dirs[0] == tmp_path / "raw" / "hits" / "gen1"
|
||||
|
||||
|
||||
def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
dirs, log_line = plan_bump_schema(
|
||||
tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01"
|
||||
)
|
||||
assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema1"]
|
||||
assert "`gen1`/`schema1`" in log_line
|
||||
|
||||
|
||||
def test_bump_schema_increments_within_its_gen(tmp_path):
|
||||
(tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True)
|
||||
(tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True)
|
||||
dirs, _ = plan_bump_schema(tmp_path, "steps", "gen1", "next schema", None, "2026-01-01")
|
||||
assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema3"]
|
||||
|
||||
|
||||
def test_bump_schema_does_not_see_other_gens_schemas(tmp_path):
|
||||
(tmp_path / "processed" / "steps" / "gen1" / "schema5").mkdir(parents=True)
|
||||
(tmp_path / "raw" / "steps" / "gen2").mkdir(parents=True)
|
||||
dirs, _ = plan_bump_schema(tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01")
|
||||
assert dirs == [tmp_path / "processed" / "steps" / "gen2" / "schema1"]
|
||||
|
||||
|
||||
def test_bump_schema_rejects_nonexistent_gen(tmp_path):
|
||||
try:
|
||||
plan_bump_schema(tmp_path, "steps", "gen9", "oops", None, "2026-01-01")
|
||||
assert False, "expected SystemExit"
|
||||
except SystemExit:
|
||||
pass
|
||||
|
||||
|
||||
def test_apply_bump_creates_dirs_and_appends_log(tmp_path):
|
||||
dirs, log_line = plan_bump_gen(tmp_path, "steps", "reason A", "alice", "2026-01-01")
|
||||
apply_bump(tmp_path, dirs, log_line)
|
||||
for d in dirs:
|
||||
assert d.is_dir()
|
||||
text = (tmp_path / "VERSIONS.md").read_text()
|
||||
assert "reason A" in text
|
||||
assert "alice" in text
|
||||
|
||||
|
||||
def test_apply_bump_appends_without_clobbering_existing_log(tmp_path):
|
||||
(tmp_path / "VERSIONS.md").write_text("# Dataset versions\n\n- existing entry\n")
|
||||
dirs, log_line = plan_bump_gen(tmp_path, "steps", "reason B", None, "2026-01-02")
|
||||
apply_bump(tmp_path, dirs, log_line)
|
||||
text = (tmp_path / "VERSIONS.md").read_text()
|
||||
assert "existing entry" in text
|
||||
assert "reason B" in text
|
||||
Reference in New Issue
Block a user