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,197 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Cut a new raw generation or processed schema version for the geant_steps
|
||||
dataset tree (see scripts/migrate_geant_steps.py for the layout):
|
||||
|
||||
raw/<kind>/<gen>/<detector>/shard-NNN.root
|
||||
processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet
|
||||
|
||||
`gen` bumps when the underlying ROOT changes (geometry/physics-list/macro).
|
||||
`schema` bumps when the parquet export (steps_to_parquet.py or similar)
|
||||
changes, and is scoped to its gen — a new gen always starts back at schema1.
|
||||
|
||||
Creates the new (empty) target directory and appends a dated, reasoned entry
|
||||
to VERSIONS.md. Defaults to a dry run; pass --execute to apply.
|
||||
|
||||
Usage:
|
||||
bump_dataset_version.py bump-gen --kind steps --reason "switched EM physics list"
|
||||
bump_dataset_version.py bump-schema --kind steps --gen gen1 --reason "added e_sec column"
|
||||
bump_dataset_version.py status
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import datetime as dt
|
||||
import re
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
GEN_RE = re.compile(r"^gen(\d+)$")
|
||||
SCHEMA_RE = re.compile(r"^schema(\d+)$")
|
||||
|
||||
|
||||
def _max_index(parent: Path, pattern: re.Pattern) -> int:
|
||||
"""Highest N across child dir names matching *pattern* (0 if none/missing)."""
|
||||
if not parent.is_dir():
|
||||
return 0
|
||||
best = 0
|
||||
for child in parent.iterdir():
|
||||
m = pattern.match(child.name)
|
||||
if m and child.is_dir():
|
||||
best = max(best, int(m.group(1)))
|
||||
return best
|
||||
|
||||
|
||||
def _git_user_name() -> str | None:
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["git", "config", "user.name"], capture_output=True, text=True, timeout=2
|
||||
)
|
||||
except OSError:
|
||||
return None
|
||||
name = out.stdout.strip()
|
||||
return name or None
|
||||
|
||||
|
||||
def plan_bump_gen(
|
||||
root: Path, kind: str, reason: str, by: str | None, date: str
|
||||
) -> tuple[list[Path], str]:
|
||||
"""New gen tag is one past the highest seen under raw/ or processed/ for *kind*
|
||||
— checking both, since a gen can exist in one tree before the other catches up."""
|
||||
next_gen = (
|
||||
max(
|
||||
_max_index(root / "raw" / kind, GEN_RE),
|
||||
_max_index(root / "processed" / kind, GEN_RE),
|
||||
)
|
||||
+ 1
|
||||
)
|
||||
gen_tag = f"gen{next_gen}"
|
||||
new_dirs = [
|
||||
root / "raw" / kind / gen_tag,
|
||||
root / "processed" / kind / gen_tag / "schema1",
|
||||
]
|
||||
by_suffix = f" ({by})" if by else ""
|
||||
log_line = f"- `{gen_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
|
||||
return new_dirs, log_line
|
||||
|
||||
|
||||
def plan_bump_schema(
|
||||
root: Path, kind: str, gen_tag: str, reason: str, by: str | None, date: str
|
||||
) -> tuple[list[Path], str]:
|
||||
if not GEN_RE.match(gen_tag):
|
||||
raise SystemExit(f"error: --gen must look like 'genN', got {gen_tag!r}")
|
||||
raw_gen_dir = root / "raw" / kind / gen_tag
|
||||
processed_gen_dir = root / "processed" / kind / gen_tag
|
||||
if not raw_gen_dir.is_dir() and not processed_gen_dir.is_dir():
|
||||
raise SystemExit(
|
||||
f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first"
|
||||
)
|
||||
next_schema = _max_index(processed_gen_dir, SCHEMA_RE) + 1
|
||||
schema_tag = f"schema{next_schema}"
|
||||
new_dirs = [processed_gen_dir / schema_tag]
|
||||
by_suffix = f" ({by})" if by else ""
|
||||
log_line = f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
|
||||
return new_dirs, log_line
|
||||
|
||||
|
||||
def apply_bump(root: Path, new_dirs: list[Path], log_line: str) -> None:
|
||||
for d in new_dirs:
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
versions_path = root / "VERSIONS.md"
|
||||
if not versions_path.exists():
|
||||
versions_path.write_text("# Dataset versions\n\n")
|
||||
with versions_path.open("a") as f:
|
||||
f.write(log_line + "\n")
|
||||
|
||||
|
||||
def print_status(root: Path) -> None:
|
||||
raw_root = root / "raw"
|
||||
if not raw_root.is_dir():
|
||||
print(f"no raw/ tree found under {root}")
|
||||
return
|
||||
for kind_dir in sorted(p for p in raw_root.iterdir() if p.is_dir()):
|
||||
kind = kind_dir.name
|
||||
gens = sorted(
|
||||
int(m.group(1))
|
||||
for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir())
|
||||
if m
|
||||
)
|
||||
print(f"kind={kind}")
|
||||
for gen in gens:
|
||||
gen_tag = f"gen{gen}"
|
||||
schema_dir = root / "processed" / kind / gen_tag
|
||||
schemas = sorted(
|
||||
int(m.group(1))
|
||||
for m in (
|
||||
SCHEMA_RE.match(p.name)
|
||||
for p in (schema_dir.iterdir() if schema_dir.is_dir() else [])
|
||||
if p.is_dir()
|
||||
)
|
||||
if m
|
||||
)
|
||||
schema_str = ", ".join(f"schema{s}" for s in schemas) or "(none)"
|
||||
print(f" {gen_tag}: {schema_str}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--root",
|
||||
default="/ceph/lbogner/geant_steps",
|
||||
help="Dataset root (default: /ceph/lbogner/geant_steps)",
|
||||
)
|
||||
sub = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
kind_help = "steps | hits | ... (default: steps)"
|
||||
|
||||
p_gen = sub.add_parser("bump-gen", help="Cut a new raw generation")
|
||||
p_gen.add_argument("--kind", default="steps", help=kind_help)
|
||||
p_gen.add_argument("--reason", required=True, help="Why this gen exists")
|
||||
p_gen.add_argument("--by", default=None, help="Attribution (default: git user.name)")
|
||||
p_gen.add_argument("--date", default=None, help="Override date (default: today, ISO)")
|
||||
p_gen.add_argument("--execute", action="store_true", help="Apply (default: dry run)")
|
||||
|
||||
p_schema = sub.add_parser("bump-schema", help="Cut a new schema within a gen")
|
||||
p_schema.add_argument("--kind", default="steps", help=kind_help)
|
||||
p_schema.add_argument("--gen", required=True, help="Existing gen tag, e.g. gen1")
|
||||
p_schema.add_argument("--reason", required=True, help="Why this schema exists")
|
||||
p_schema.add_argument("--by", default=None, help="Attribution (default: git user.name)")
|
||||
p_schema.add_argument("--date", default=None, help="Override date (default: today, ISO)")
|
||||
p_schema.add_argument("--execute", action="store_true", help="Apply (default: dry run)")
|
||||
|
||||
sub.add_parser("status", help="List existing gens/schemas per kind")
|
||||
|
||||
args = parser.parse_args()
|
||||
root = Path(args.root)
|
||||
if not root.is_dir():
|
||||
parser.error(f"{root} is not a directory")
|
||||
|
||||
if args.command == "status":
|
||||
print_status(root)
|
||||
return
|
||||
|
||||
date = args.date or dt.date.today().isoformat()
|
||||
by = args.by if args.by is not None else _git_user_name()
|
||||
|
||||
if args.command == "bump-gen":
|
||||
new_dirs, log_line = plan_bump_gen(root, args.kind, args.reason, by, date)
|
||||
else:
|
||||
new_dirs, log_line = plan_bump_schema(
|
||||
root, args.kind, args.gen, args.reason, by, date
|
||||
)
|
||||
|
||||
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===")
|
||||
print("new directories:")
|
||||
for d in new_dirs:
|
||||
print(f" {d}")
|
||||
print("VERSIONS.md entry:")
|
||||
print(f" {log_line}")
|
||||
|
||||
if not args.execute:
|
||||
print("\nDry run only — pass --execute to apply.")
|
||||
return
|
||||
|
||||
apply_bump(root, new_dirs, log_line)
|
||||
print("\nDone.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user