lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge energy-conservation-poc into phase2-secondary-prediction

Browse Source 
Brings the energy-conservation PoC work (dwarf CLI unification, dwarf
status improvements, predict --comment, ODE-step comparison scripts,
predict-parquet-only analysis refactor) onto the Phase 2 branch.

Conflict resolution:
- giant/analysis.py: took the energy-conservation-poc version wholesale.
  That branch deliberately removed the live checkpoint+sampler diagnostics
  path (ModelBundle/load_model_bundle/make_val_loader/collect_samples) in
  favor of reading `giant predict --coord local` parquet output. Phase 2's
  only edits to this file adapted the removed path to the new dataset API,
  so nothing Phase-2-specific is lost; no external code called those funcs.

Fixes for pre-existing breakage surfaced by the merge (both predate it):
- giant/cli.py: predict's `_process` unpacked build_features into 5 values,
  but Phase 2 made it return 8 (added n_sec/sec_cont/sec_pdg_idx). Expanded
  the unpack; `giant predict --coord local` would have crashed otherwise.
- tests/test_steps_to_parquet.py: Phase 2 renamed _add_secondary_energy ->
  _add_secondary_attributes without updating this test. Renamed the calls
  and extended the fixture with the pdg/pre_d{x,y,z} columns the expanded
  function reads; e_sec assertions unchanged.
- analysis/compare_ode_steps_energy_conservation.py: E731 lambda assignment
  (added in the un-linted final PoC commit) rewritten as a def.

ruff, ty, and pytest (179 passed) all green.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-06 12:18:30 +02:00
parent b25c6967ab 1a22ae022c
commit 4ba419ebe4
33 changed files with 2316 additions and 964 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/bump_dataset_version.py
+526 -223
View File
@@ -1,4 +1,3 @@
#!/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):
@@ -7,30 +6,67 @@ dataset tree (see scripts/migrate_geant_steps.py for the layout):
pools/<detector>/<pool>.manifest
`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.
`schema` bumps when the parquet export (`dwarf convert` 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.
to VERSIONS.md. Defaults to a dry run; pass execute=True 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 update-manifest pools/pbwo4/full.manifest [--schema schema2]
bump_dataset_version.py create-manifest --output pools/pbwo4/train.manifest a.parquet b.parquet
bump_dataset_version.py status
See `uv run dwarf bump-gen/bump-schema/update-manifest/create-manifest/status
--help` for the CLI.
"""
import argparse
import datetime as dt
import os
import re
import subprocess
import sys
from pathlib import Path
GEN_RE = re.compile(r"^gen(\d+)$")
SCHEMA_RE = re.compile(r"^schema(\d+)$")
# Match the log lines written by apply_bump()/plan_bump_gen()/plan_bump_schema():
# - `gen2` (kind=steps) — 2026-01-01 — reason text (by)
# - `gen2`/`schema2` (kind=steps) — 2026-01-01 — reason text (by)
VERSIONS_SCHEMA_LINE_RE = re.compile(
r"^- `(?P<gen>gen\d+)`/`(?P<schema>schema\d+)` \(kind=(?P<kind>[\w-]+)\)"
r" — \d{4}-\d{2}-\d{2} — (?P<reason>.+)$"
)
VERSIONS_GEN_LINE_RE = re.compile(
r"^- `(?P<gen>gen\d+)` \(kind=(?P<kind>[\w-]+)\) — \d{4}-\d{2}-\d{2} — (?P<reason>.+)$"
)
# One color per tree level in `dwarf status` output, so the eye can jump
# straight to e.g. "all the schema rows" or "all the totals".
_LEVEL_COLORS = {
"kind": "\033[1;36m", # bold cyan — top-level kind/ header
"gen": "\033[1;33m", # bold yellow — genN row + kind total
"bucket": "\033[34m", # blue — raw/processed subtotals
"schema": "\033[32m", # green — schemaN rows
"root": "\033[1;35m", # bold magenta — derived/, pools/, grand total
"reason": "\033[2m", # dim — VERSIONS.md reason extract under gen/schema rows
}
_RESET = "\033[0m"
def _use_color() -> bool:
return sys.stdout.isatty() and os.environ.get("NO_COLOR") is None
def _colorize(text: str, level: str) -> str:
if not _use_color():
return text
return f"{_LEVEL_COLORS[level]}{text}{_RESET}"
def _find_component_idx(abs_path: Path, pattern: re.Pattern) -> int | None:
"""Return the index of the first path component matching *pattern*, or None."""
for i, part in enumerate(abs_path.parts):
if pattern.match(part):
return i
return None
def _max_index(parent: Path, pattern: re.Pattern) -> int:
"""Highest N across child dir names matching *pattern* (0 if none/missing)."""
@@ -56,18 +92,28 @@ def _git_user_name() -> str | None:
def plan_bump_gen(
root: Path, kind: str, reason: str, by: str | None, date: str
root: Path,
kind: str,
reason: str,
by: str | None,
date: str,
target: str | None = None,
) -> 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),
"""New gen tag is one past the highest seen under raw/ or processed/ for *kind*,
or *target* if explicitly provided."""
if target is not None:
if not GEN_RE.match(target):
raise SystemExit(f"error: --to must look like 'genN', got {target!r}")
gen_tag = target
else:
next_gen = (
max(
_max_index(root / "raw" / kind, GEN_RE),
_max_index(root / "processed" / kind, GEN_RE),
)
+ 1
)
+ 1
)
gen_tag = f"gen{next_gen}"
gen_tag = f"gen{next_gen}"
new_dirs = [
root / "raw" / kind / gen_tag,
root / "processed" / kind / gen_tag / "schema1",
@@ -78,7 +124,13 @@ def plan_bump_gen(
def plan_bump_schema(
root: Path, kind: str, gen_tag: str, reason: str, by: str | None, date: str
root: Path,
kind: str,
gen_tag: str,
reason: str,
by: str | None,
date: str,
target: str | None = None,
) -> tuple[list[Path], str]:
if not GEN_RE.match(gen_tag):
raise SystemExit(f"error: --gen must look like 'genN', got {gen_tag!r}")
@@ -88,11 +140,18 @@ def plan_bump_schema(
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}"
if target is not None:
if not SCHEMA_RE.match(target):
raise SystemExit(f"error: --to must look like 'schemaN', got {target!r}")
schema_tag = target
else:
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}"
log_line = (
f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
)
return new_dirs, log_line
@@ -106,11 +165,176 @@ def apply_bump(root: Path, new_dirs: list[Path], log_line: str) -> None:
f.write(log_line + "\n")
def _du(path: Path) -> int:
"""Total size in bytes of all regular files under *path* (0 if missing)."""
if not path.is_dir():
return 0
total = 0
for dirpath, _dirnames, filenames in os.walk(path):
for name in filenames:
fp = Path(dirpath) / name
try:
total += fp.stat().st_size
except OSError:
pass
return total
def _count_files(path: Path) -> int:
"""Total number of regular files under *path*, recursively (0 if missing)."""
if not path.is_dir():
return 0
total = 0
for _dirpath, _dirnames, filenames in os.walk(path):
total += len(filenames)
return total
# Must match giant.data.loader.MANIFEST_SUFFIX.
MANIFEST_SUFFIX = ".manifest"
def _manifest_referenced_files(pools_root: Path) -> set[Path]:
"""Resolved absolute paths of every file listed in any *.manifest under *pools_root*."""
referenced: set[Path] = set()
if not pools_root.is_dir():
return referenced
for manifest_path in pools_root.rglob(f"*{MANIFEST_SUFFIX}"):
try:
referenced.update(_resolve_manifest_files(manifest_path))
except OSError:
continue
return referenced
def _referenced_root_count(
raw_gen_dir: Path, processed_gen_dir: Path
) -> tuple[int, int]:
"""(total .root files, count with a same-named .parquet under any schema) for one gen."""
if not raw_gen_dir.is_dir():
return 0, 0
parquet_stems: set[tuple[str, str]] = set()
if processed_gen_dir.is_dir():
for schema_dir in processed_gen_dir.iterdir():
if not schema_dir.is_dir():
continue
for detector_dir in schema_dir.iterdir():
if not detector_dir.is_dir():
continue
for f in detector_dir.iterdir():
if f.is_file() and f.suffix == ".parquet":
parquet_stems.add((detector_dir.name, f.stem))
total = 0
referenced = 0
for detector_dir in raw_gen_dir.iterdir():
if not detector_dir.is_dir():
continue
for f in detector_dir.iterdir():
if f.is_file() and f.suffix == ".root":
total += 1
if (detector_dir.name, f.stem) in parquet_stems:
referenced += 1
return total, referenced
def _referenced_parquet_count(
schema_dir: Path, manifest_referenced: set[Path]
) -> tuple[int, int]:
"""(total .parquet files, count listed in at least one manifest) for one schema dir."""
if not schema_dir.is_dir():
return 0, 0
total = 0
referenced = 0
for detector_dir in schema_dir.iterdir():
if not detector_dir.is_dir():
continue
for f in detector_dir.iterdir():
if f.is_file() and f.suffix == ".parquet":
total += 1
if f.resolve() in manifest_referenced:
referenced += 1
return total, referenced
def _parse_versions(
versions_path: Path,
) -> tuple[dict[tuple[str, str], str], dict[tuple[str, str, str], str]]:
"""Read VERSIONS.md and return (gen_reasons, schema_reasons) keyed by
(kind, gen_tag) and (kind, gen_tag, schema_tag) respectively. Later entries
for the same key win, since VERSIONS.md is append-only and chronological."""
gen_reasons: dict[tuple[str, str], str] = {}
schema_reasons: dict[tuple[str, str, str], str] = {}
if not versions_path.is_file():
return gen_reasons, schema_reasons
for line in versions_path.read_text().splitlines():
line = line.strip()
m = VERSIONS_SCHEMA_LINE_RE.match(line)
if m:
schema_reasons[(m["kind"], m["gen"], m["schema"])] = m["reason"]
continue
m = VERSIONS_GEN_LINE_RE.match(line)
if m:
gen_reasons[(m["kind"], m["gen"])] = m["reason"]
return gen_reasons, schema_reasons
def _truncate(text: str, width: int = 72) -> str:
text = text.strip()
if len(text) <= width:
return text
return text[: width - 1].rstrip() + "…"
def _human_size(n: int) -> str:
size = float(n)
for unit in ("B", "KB", "MB", "GB", "TB"):
if size < 1024 or unit == "TB":
return f"{size:.1f} {unit}" if unit != "B" else f"{int(size)} {unit}"
size /= 1024
return f"{size:.1f} TB"
_LABEL_WIDTH = 26
_COUNT_WIDTH = 20
_ROW_WIDTH = _LABEL_WIDTH + _COUNT_WIDTH + 10
def _count_str(count: int, referenced: int | None = None) -> str:
files = "file" if count == 1 else "files"
if referenced is not None:
return f"{count} {files} ({referenced} ref)"
return f"{count} {files}"
def _row(
label: str,
size_bytes: int,
indent: int = 0,
level: str | None = None,
count: int | None = None,
referenced: int | None = None,
) -> str:
text = " " * indent + label
count_str = _count_str(count, referenced) if count is not None else ""
size_str = _human_size(size_bytes)
row = f"{text:<{_LABEL_WIDTH}}{count_str:<{_COUNT_WIDTH}}{size_str:>10}"
return _colorize(row, level) if level else row
def _reason_line(reason: str, indent: int) -> str:
return _colorize(" " * indent + "↳ " + _truncate(reason), "reason")
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
manifest_referenced = _manifest_referenced_files(root / "pools")
gen_reasons, schema_reasons = _parse_versions(root / "VERSIONS.md")
grand_total = 0
grand_files = 0
for kind_dir in sorted(p for p in raw_root.iterdir() if p.is_dir()):
kind = kind_dir.name
gens = sorted(
@@ -118,10 +342,15 @@ def print_status(root: Path) -> None:
for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir())
if m
)
print(f"kind={kind}")
print(_colorize(f"{kind}/", "kind"))
kind_total = 0
kind_files = 0
for gen in gens:
gen_tag = f"gen{gen}"
schema_dir = root / "processed" / kind / gen_tag
raw_gen_dir = raw_root / kind / gen_tag
raw_size = _du(raw_gen_dir)
processed_gen_dir = root / "processed" / kind / gen_tag
schema_dir = processed_gen_dir
schemas = sorted(
int(m.group(1))
for m in (
@@ -131,26 +360,98 @@ def print_status(root: Path) -> None:
)
if m
)
schema_str = ", ".join(f"schema{s}" for s in schemas) or "(none)"
print(f" {gen_tag}: {schema_str}")
schema_sizes = {s: _du(schema_dir / f"schema{s}") for s in schemas}
schema_counts = {
s: _referenced_parquet_count(
schema_dir / f"schema{s}", manifest_referenced
)
for s in schemas
}
processed_size = sum(schema_sizes.values())
processed_files = sum(c[0] for c in schema_counts.values())
processed_referenced = sum(c[1] for c in schema_counts.values())
raw_files, raw_referenced = _referenced_root_count(
raw_gen_dir, processed_gen_dir
)
gen_total = raw_size + processed_size
gen_files = raw_files + processed_files
kind_total += gen_total
kind_files += gen_files
print(_row(gen_tag, gen_total, indent=1, level="gen", count=gen_files))
gen_reason = gen_reasons.get((kind, gen_tag))
if gen_reason:
print(_reason_line(gen_reason, indent=2))
print(
_row(
"raw",
raw_size,
indent=2,
level="bucket",
count=raw_files,
referenced=raw_referenced,
)
)
print(
_row(
"processed",
processed_size,
indent=2,
level="bucket",
count=processed_files,
referenced=processed_referenced,
)
)
if schemas:
for s in schemas:
s_total, s_referenced = schema_counts[s]
print(
_row(
f"schema{s}",
schema_sizes[s],
indent=3,
level="schema",
count=s_total,
referenced=s_referenced,
)
)
schema_reason = schema_reasons.get((kind, gen_tag, f"schema{s}"))
if schema_reason:
print(_reason_line(schema_reason, indent=4))
else:
print(_colorize(" (none)", "schema"))
print(
_row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files)
)
print()
grand_total += kind_total
grand_files += kind_files
derived_dir = root / "derived"
pools_dir = root / "pools"
derived_size = _du(derived_dir)
pools_size = _du(pools_dir)
derived_files = _count_files(derived_dir)
pools_files = _count_files(pools_dir)
grand_total += derived_size + pools_size
grand_files += derived_files + pools_files
print(_row("derived/", derived_size, level="root", count=derived_files))
print(_row("pools/", pools_size, level="root", count=pools_files))
print("-" * _ROW_WIDTH)
print(_row("grand total", grand_total, level="root", count=grand_files))
# ---------------------------------------------------------------------------
# update-manifest
# ---------------------------------------------------------------------------
def _find_schema_idx(abs_path: Path) -> int | None:
"""Return the index of the first schemaN component in abs_path.parts, or None."""
for i, part in enumerate(abs_path.parts):
if SCHEMA_RE.match(part):
return i
return None
def plan_update_manifest(
manifest_path: Path, target_schema: str | None
manifest_path: Path,
target_schema: str | None,
target_gen: str | None = None,
) -> tuple[list[tuple[str, str | None]], list[Path]]:
"""Parse a manifest and plan schema replacements for each data line.
"""Parse a manifest and plan gen/schema replacements for each data line.
Returns:
lines: list of (original_line, new_relative_path_or_None)
@@ -174,42 +475,59 @@ def plan_update_manifest(
continue
old_abs = (manifest_dir / stripped).resolve()
schema_idx = _find_schema_idx(old_abs)
if schema_idx is None:
result.append((raw, None))
continue
parts = list(old_abs.parts)
old_schema = parts[schema_idx]
gen_dir = Path(*parts[:schema_idx])
changed = False
if target_schema is not None:
new_schema = target_schema
else:
if gen_dir not in schema_cache:
n = _max_index(gen_dir, SCHEMA_RE)
if n == 0:
raise SystemExit(f"error: no schema dirs found under {gen_dir}")
schema_cache[gen_dir] = f"schema{n}"
new_schema = schema_cache[gen_dir]
# Apply gen replacement first (shifts subsequent indices).
if target_gen is not None:
gen_idx = _find_component_idx(Path(*parts), GEN_RE)
if gen_idx is not None and parts[gen_idx] != target_gen:
parts[gen_idx] = target_gen
changed = True
if new_schema == old_schema:
result.append((raw, None))
continue
schema_idx = _find_component_idx(Path(*parts), SCHEMA_RE)
if schema_idx is not None:
old_schema = parts[schema_idx]
gen_dir = Path(*parts[:schema_idx])
parts[schema_idx] = new_schema
if target_schema is not None:
new_schema = target_schema
else:
if gen_dir not in schema_cache:
n = _max_index(gen_dir, SCHEMA_RE)
if n == 0:
raise SystemExit(f"error: no schema dirs found under {gen_dir}")
schema_cache[gen_dir] = f"schema{n}"
new_schema = schema_cache[gen_dir]
if new_schema != old_schema:
parts[schema_idx] = new_schema
changed = True
# Check existence for every data line, not just ones whose gen/schema
# actually changed — an already-correct-looking line can still point
# at a file that was deleted or moved out-of-band.
new_abs = Path(*parts)
if not new_abs.exists():
missing.append(new_abs)
if not changed:
result.append((raw, None))
continue
new_rel = os.path.relpath(new_abs, start=manifest_dir)
result.append((raw, new_rel))
return result, missing
def apply_update_manifest(manifest_path: Path, lines: list[tuple[str, str | None]]) -> None:
out = [replacement if replacement is not None else original for original, replacement in lines]
def apply_update_manifest(
manifest_path: Path, lines: list[tuple[str, str | None]]
) -> None:
out = [
replacement if replacement is not None else original
for original, replacement in lines
]
manifest_path.write_text("\n".join(out) + "\n")
@@ -217,6 +535,7 @@ def apply_update_manifest(manifest_path: Path, lines: list[tuple[str, str | None
# create-manifest
# ---------------------------------------------------------------------------
def _resolve_manifest_files(manifest_path: Path) -> list[Path]:
"""Read a manifest and return its entries as resolved absolute paths."""
files = []
@@ -287,183 +606,167 @@ def apply_create_manifest(output_path: Path, lines: list[str]) -> None:
# ---------------------------------------------------------------------------
# CLI
# CLI entry points (called from scripts/dwarf.py)
# ---------------------------------------------------------------------------
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--root",
default="/ceph/lbogner/geant_steps",
help="Dataset root for bump-gen/bump-schema/status (default: /ceph/lbogner/geant_steps)",
)
sub = parser.add_subparsers(dest="command", required=True)
kind_help = "steps | hits | ... (default: steps)"
def run_status(root: str) -> None:
root_path = Path(root)
if not root_path.is_dir():
raise SystemExit(f"error: {root_path} is not a directory")
print_status(root_path)
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)")
def _run_bump(
kind: str,
reason: str,
by: str | None,
date: str | None,
execute: bool,
root: str,
gen: str | None,
to: str | None,
) -> None:
root_path = Path(root)
if not root_path.is_dir():
raise SystemExit(f"error: {root_path} is not a directory")
p_update = sub.add_parser(
"update-manifest",
help="Repoint manifest(s) to a new schema, verifying all target files exist",
)
p_update.add_argument(
"manifests", nargs="+", metavar="MANIFEST", help="One or more .manifest files to update"
)
p_update.add_argument(
"--schema",
default=None,
metavar="schemaN",
help="Target schema tag (default: highest schema found in the same gen dir)",
)
p_update.add_argument("--execute", action="store_true", help="Write updated manifests (default: dry run)")
date = date or dt.date.today().isoformat()
by = by if by is not None else _git_user_name()
if gen is None:
new_dirs, log_line = plan_bump_gen(root_path, kind, reason, by, date, to)
else:
new_dirs, log_line = plan_bump_schema(
root_path, kind, gen, reason, by, date, to
)
p_create = sub.add_parser(
"create-manifest",
help="Create a new manifest from a list of parquet files",
)
dest_group = p_create.add_mutually_exclusive_group(required=True)
dest_group.add_argument(
"--output", "-o", metavar="PATH", help="Explicit path for the new .manifest file"
)
dest_group.add_argument(
"--pool", metavar="DETECTOR",
help="Detector name; combined with --type and --root to form <root>/pools/<detector>/<type>.manifest",
)
p_create.add_argument(
"--type", choices=["full", "holdout", "dev"],
help="Pool type — full, holdout, or dev (required with --pool)",
)
p_create.add_argument(
"files", nargs="+", metavar="FILE", help="Parquet files to include"
)
p_create.add_argument("--execute", action="store_true", help="Write the manifest (default: dry run)")
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
print("new directories:")
for d in new_dirs:
print(f" {d}")
print("VERSIONS.md entry:")
print(f" {log_line}")
sub.add_parser("status", help="List existing gens/schemas per kind")
if not execute:
print("\nDry run only — pass --execute to apply.")
return
apply_bump(root_path, new_dirs, log_line)
print("\nDone.")
args = parser.parse_args()
# Commands that need --root
if args.command in ("bump-gen", "bump-schema", "status"):
root = Path(args.root)
if not root.is_dir():
parser.error(f"{root} is not a directory")
def run_bump_gen(
kind: str,
reason: str,
by: str | None,
date: str | None,
execute: bool,
root: str,
to: str | None = None,
) -> None:
_run_bump(kind, reason, by, date, execute, root, gen=None, to=to)
if args.command == "status":
print_status(root)
def run_bump_schema(
kind: str,
gen: str,
reason: str,
by: str | None,
date: str | None,
execute: bool,
root: str,
to: str | None = None,
) -> None:
_run_bump(kind, reason, by, date, execute, root, gen=gen, to=to)
def run_update_manifest(
manifests: list[str],
schema: str | None,
execute: bool,
gen: str | None = None,
) -> None:
if schema and not SCHEMA_RE.match(schema):
raise SystemExit(f"error: --schema must look like 'schemaN', got {schema!r}")
if gen and not GEN_RE.match(gen):
raise SystemExit(f"error: --gen must look like 'genN', got {gen!r}")
all_plans: list[tuple[Path, list[tuple[str, str | None]]]] = []
all_missing: list[Path] = []
for raw in manifests:
mp = Path(raw)
plan, missing = plan_update_manifest(mp, schema, gen)
all_plans.append((mp.resolve(), plan))
all_missing.extend(missing)
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
for mp, plan in all_plans:
changes = [(old, new) for old, new in plan if new is not None]
print(f"\n{mp} ({len(changes)} path(s) to update)")
for old, new in changes:
print(f" - {old.strip()}")
print(f" + {new}")
if all_missing:
print(f"\nMISSING ({len(all_missing)} file(s) — target paths do not exist):")
for p in all_missing:
print(f" {p}")
if execute:
raise SystemExit("error: refusing to write manifests with missing targets")
if not execute:
print("\nDry run only — pass --execute to apply.")
return
if args.command in ("bump-gen", "bump-schema"):
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)
for mp, plan in all_plans:
apply_update_manifest(mp, plan)
print("\nDone.")
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.")
def run_create_manifest(
files: list[str],
execute: bool,
output: str | None = None,
pool: str | None = None,
type_: str | None = None,
root: str = "/ceph/lbogner/geant_steps",
) -> None:
if (output is None) == (pool is None):
raise SystemExit("error: exactly one of --output or --pool is required")
if pool is not None and type_ is None:
raise SystemExit("error: --type is required when --pool is given")
if pool is not None:
output_path = Path(root) / "pools" / pool / f"{type_}.manifest"
else:
assert output is not None # guaranteed by the exclusivity check above
output_path = Path(output)
parquet_files = [Path(f) for f in files]
lines, missing, resolved = plan_create_manifest(output_path, parquet_files)
overlaps = check_holdout_overlap(output_path, resolved)
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
print(f"manifest: {output_path.resolve()}")
for line in lines:
print(f" {line}")
if missing:
print(f"\nMISSING ({len(missing)} file(s) do not exist):")
for p in missing:
print(f" {p}")
if overlaps:
print(f"\nHOLDOUT OVERLAP ({len(overlaps)} file(s) appear in other manifests):")
for name, f in overlaps:
print(f" {f} (also in {name})")
if (missing or overlaps) and execute:
raise SystemExit("error: refusing to write manifest (see above)")
if not execute:
print("\nDry run only — pass --execute to apply.")
return
if args.command == "update-manifest":
if args.schema and not SCHEMA_RE.match(args.schema):
parser.error(f"--schema must look like 'schemaN', got {args.schema!r}")
all_plans: list[tuple[Path, list[tuple[str, str | None]]]] = []
all_missing: list[Path] = []
for raw in args.manifests:
mp = Path(raw)
plan, missing = plan_update_manifest(mp, args.schema)
all_plans.append((mp.resolve(), plan))
all_missing.extend(missing)
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===")
for mp, plan in all_plans:
changes = [(old, new) for old, new in plan if new is not None]
print(f"\n{mp} ({len(changes)} path(s) to update)")
for old, new in changes:
print(f" - {old.strip()}")
print(f" + {new}")
if all_missing:
print(f"\nMISSING ({len(all_missing)} file(s) — target paths do not exist):")
for p in all_missing:
print(f" {p}")
if args.execute:
raise SystemExit("error: refusing to write manifests with missing targets")
if not args.execute:
print("\nDry run only — pass --execute to apply.")
return
for mp, plan in all_plans:
apply_update_manifest(mp, plan)
print("\nDone.")
return
if args.command == "create-manifest":
if args.pool is not None and args.type is None:
parser.error("--type is required when --pool is given")
if args.pool is not None:
root = Path(args.root)
output_path = root / "pools" / args.pool / f"{args.type}.manifest"
else:
output_path = Path(args.output)
parquet_files = [Path(f) for f in args.files]
lines, missing, resolved = plan_create_manifest(output_path, parquet_files)
overlaps = check_holdout_overlap(output_path, resolved)
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===")
print(f"manifest: {output_path.resolve()}")
for line in lines:
print(f" {line}")
if missing:
print(f"\nMISSING ({len(missing)} file(s) do not exist):")
for p in missing:
print(f" {p}")
if overlaps:
print(f"\nHOLDOUT OVERLAP ({len(overlaps)} file(s) appear in other manifests):")
for name, f in overlaps:
print(f" {f} (also in {name})")
if (missing or overlaps) and args.execute:
raise SystemExit("error: refusing to write manifest (see above)")
if not args.execute:
print("\nDry run only — pass --execute to apply.")
return
apply_create_manifest(output_path, lines)
print("\nDone.")
if __name__ == "__main__":
main()
apply_create_manifest(output_path, lines)
print("\nDone.")
scripts/create_root_files.py
+94 -88
View File
@@ -1,4 +1,3 @@
#!/usr/bin/env python3
"""Generate new ROOT shards by running a minicalosim executable (e.g.
run_pbwo4, run_sampling) and filing the output into the dataset's raw/ tree:
@@ -15,19 +14,11 @@ appear there, and moves it to the next free shard index for that detector
(existing shards are never overwritten).
--gen must already exist under raw/<kind>/ — create one first with
bump_dataset_version.py bump-gen.
`dwarf bump-gen`.
Usage:
create_root_files.py --executable build/run_pbwo4 --detector pbwo4 \\
--gen gen1 --num-files 4 --events-per-file 10000 --execute
create_root_files.py --executable build/run_sampling \\
--detector sampling_pb_scint:pb_scint \\
--detector sampling_fe_scint:fe_scint \\
--gen gen1 --num-files 4 --events-per-file 10000 --jobs 8 --execute
See `uv run dwarf make-root --help` for the CLI.
"""
import argparse
import os
import re
import shutil
@@ -70,7 +61,9 @@ def parse_detector_spec(spec: str) -> tuple[str, str | None]:
if ":" in spec:
label, config = spec.split(":", 1)
if not label or not config:
raise PlanError(f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG")
raise PlanError(
f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG"
)
return label, config
return spec, None
@@ -120,7 +113,10 @@ def run_job(
gen: str,
tmp_root: Path,
) -> JobResult:
workdir = tmp_root / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
workdir = (
tmp_root
/ f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
)
workdir.mkdir(parents=True)
cmd = [str(executable)]
@@ -132,25 +128,42 @@ def run_job(
if result.returncode != 0:
return JobResult(
job, False, None,
job,
False,
None,
f"executable exited {result.returncode}",
result.stdout, result.stderr,
result.stdout,
result.stderr,
)
produced = sorted(workdir.glob("*.root"))
if len(produced) != 1:
return JobResult(
job, False, None,
job,
False,
None,
f"expected exactly one .root output in {workdir}, found {len(produced)}: "
f"{[p.name for p in produced]}",
result.stdout, result.stderr,
result.stdout,
result.stderr,
)
dest = dataset_root / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
dest = (
dataset_root
/ "raw"
/ kind
/ gen
/ job.detector
/ f"shard-{job.shard_index:03d}.root"
)
if dest.exists():
return JobResult(
job, False, None, f"refusing to overwrite existing {dest}",
result.stdout, result.stderr,
job,
False,
None,
f"refusing to overwrite existing {dest}",
result.stdout,
result.stderr,
)
dest.parent.mkdir(parents=True, exist_ok=True)
@@ -175,7 +188,14 @@ def run_all(
with ThreadPoolExecutor(max_workers=max_workers) as pool:
futures = {
pool.submit(
run_job, job, executable, events_per_file, dataset_root, kind, gen, tmp_root
run_job,
job,
executable,
events_per_file,
dataset_root,
kind,
gen,
tmp_root,
): job
for job in jobs
}
@@ -192,78 +212,65 @@ def run_all(
return results
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--executable", required=True, type=Path, help="Built minicalosim run_* executable"
)
parser.add_argument(
"--detector",
action="append",
required=True,
metavar="NAME[:CONFIG]",
help="Dataset detector label, optionally with ':CONFIG' to pass as the "
"executable's config-name argument (e.g. sampling_pb_scint:pb_scint). "
"Omit ':CONFIG' for executables that take no config selector (e.g. run_pbwo4). "
"Repeatable.",
)
parser.add_argument(
"--num-files", type=int, required=True, help="New shards to create per detector"
)
parser.add_argument(
"--events-per-file", type=int, required=True, help="nEvents passed to the executable"
)
parser.add_argument("--kind", default="steps", help="steps | hits | ... (default: steps)")
parser.add_argument("--gen", required=True, help="Existing gen tag under raw/<kind>/, e.g. gen1")
parser.add_argument(
"--dataset-root",
default="/ceph/lbogner/geant_steps",
help="Dataset root (default: /ceph/lbogner/geant_steps)",
)
parser.add_argument(
"-j", "--jobs", type=int, default=4, help="Parallel simulation runs (default: 4)"
)
parser.add_argument(
"--execute", action="store_true", help="Actually run jobs (default: dry run / print plan)"
)
return parser
def run_make_root(
executable: Path,
detector: list[str],
num_files: int,
events_per_file: int,
kind: str,
gen: str,
dataset_root: str,
jobs: int,
execute: bool,
) -> None:
if jobs < 1:
raise SystemExit("error: --jobs must be >= 1")
if num_files < 1:
raise SystemExit("error: --num-files must be >= 1")
if events_per_file < 1:
raise SystemExit("error: --events-per-file must be >= 1")
if not executable.is_file() or not os.access(executable, os.X_OK):
raise SystemExit(f"error: {executable} is not an executable file")
def main() -> None:
parser = build_parser()
args = parser.parse_args()
if args.jobs < 1:
parser.error("--jobs must be >= 1")
if args.num_files < 1:
parser.error("--num-files must be >= 1")
if args.events_per_file < 1:
parser.error("--events-per-file must be >= 1")
if not args.executable.is_file() or not os.access(args.executable, os.X_OK):
parser.error(f"{args.executable} is not an executable file")
dataset_root = Path(args.dataset_root)
dataset_root_path = Path(dataset_root)
try:
jobs = plan_jobs(args.detector, args.num_files, dataset_root, args.kind, args.gen)
planned_jobs = plan_jobs(detector, num_files, dataset_root_path, kind, gen)
except PlanError as exc:
parser.error(str(exc))
raise SystemExit(f"error: {exc}")
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===")
print(f"executable: {args.executable}")
for job in jobs:
cmd = [str(args.executable)] + ([job.config] if job.config else []) + [str(args.events_per_file)]
dest = dataset_root / "raw" / args.kind / args.gen / job.detector / f"shard-{job.shard_index:03d}.root"
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
print(f"executable: {executable}")
for job in planned_jobs:
cmd = (
[str(executable)]
+ ([job.config] if job.config else [])
+ [str(events_per_file)]
)
dest = (
dataset_root_path
/ "raw"
/ kind
/ gen
/ job.detector
/ f"shard-{job.shard_index:03d}.root"
)
print(f" {' '.join(cmd)} -> {dest}")
if not args.execute:
if not execute:
print("\nDry run only — pass --execute to apply.")
return
tmp_root = dataset_root / ".sim-tmp"
tmp_root = dataset_root_path / ".sim-tmp"
tmp_root.mkdir(parents=True, exist_ok=True)
results = run_all(
jobs, args.executable, args.events_per_file, dataset_root, args.kind, args.gen,
max_workers=args.jobs, tmp_root=tmp_root,
planned_jobs,
executable,
events_per_file,
dataset_root_path,
kind,
gen,
max_workers=jobs,
tmp_root=tmp_root,
)
if tmp_root.is_dir() and not any(tmp_root.iterdir()):
tmp_root.rmdir()
@@ -272,11 +279,10 @@ def main() -> None:
if failures:
print(f"\n{len(failures)} of {len(results)} job(s) failed:", file=sys.stderr)
for r in failures:
print(f" {r.job.detector} shard-{r.job.shard_index:03d}: {r.message}", file=sys.stderr)
sys.exit(1)
print(
f" {r.job.detector} shard-{r.job.shard_index:03d}: {r.message}",
file=sys.stderr,
)
raise SystemExit(1)
print(f"\nAll {len(results)} job(s) completed.")
if __name__ == "__main__":
main()
scripts/dwarf.py
+400
View File
@@ -0,0 +1,400 @@
"""dwarf — little helper to `giant`: dataset/tooling CLI for the geant_steps pipeline.
Unifies the standalone scripts/*.py conversion, migration, versioning, and
simulation-fanout tools into one Typer app so there's a single command name
(and `--help`) to remember instead of five differently-hyphenated ones.
"""
from enum import Enum
from pathlib import Path
from typing import Optional
import typer
from typing_extensions import Annotated
from scripts.bump_dataset_version import (
run_bump_gen,
run_bump_schema,
run_create_manifest,
run_status,
run_update_manifest,
)
from scripts.create_root_files import run_make_root
from scripts.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
from scripts.migrate_geant_steps import run_migration
from scripts.steps_to_parquet import convert_steps_to_parquet
from scripts.steps_to_parquet_parallel import run_parallel_job
app = typer.Typer(no_args_is_help=True)
_DATASET_ROOT_DEFAULT = Path("/ceph/lbogner/geant_steps")
@app.callback()
def _main() -> None:
"""dwarf — dataset/tooling CLI (ROOT<->parquet conversion, dataset versioning, sim fanout)."""
class Compression(str, Enum):
snappy = "snappy"
lz4 = "lz4"
zstd = "zstd"
gzip = "gzip"
none = "none"
class PoolType(str, Enum):
full = "full"
holdout = "holdout"
dev = "dev"
@app.command()
def convert(
root_files: Annotated[list[Path], typer.Argument(help="Input ROOT file(s)")],
output: Annotated[
Optional[Path],
typer.Option(
"--output",
"-o",
help="Output Parquet file (default: <input>.parquet). Only valid "
"with a single input file and --jobs 1.",
),
] = None,
batch_size: Annotated[
str,
typer.Option(
"--batch-size",
help="Uproot read batch size, e.g. '100 MB' or '500000' (rows)",
),
] = "100 MB",
tree: Annotated[
str, typer.Option("--tree", help="Tree name inside the ROOT file")
] = "Steps",
compression: Annotated[
Compression, typer.Option("--compression", help="Parquet compression codec")
] = Compression.snappy,
jobs: Annotated[
int,
typer.Option(
"--jobs",
"-j",
help="Convert N files in parallel, resolving each destination from "
"--dataset-root/--schema (default: 1, sequential, any file layout)",
),
] = 1,
dataset_root: Annotated[
Path,
typer.Option(
"--dataset-root",
help="Dataset root containing raw/ and processed/ (only used with --jobs > 1)",
),
] = _DATASET_ROOT_DEFAULT,
schema: Annotated[
Optional[str],
typer.Option(
"--schema",
help="Schema tag to write parquets under, e.g. schema2 (only used "
"with --jobs > 1; default: highest schemaN already under "
"processed/<kind>/<gen>/)",
),
] = None,
) -> None:
"""Convert ROOT Steps tree(s) to Parquet."""
if jobs < 1:
typer.echo("error: --jobs must be >= 1", err=True)
raise typer.Exit(1)
compression_value = (
"uncompressed" if compression is Compression.none else compression.value
)
if jobs == 1:
if output is not None and len(root_files) > 1:
typer.echo(
"error: --output can only be used with a single input file", err=True
)
raise typer.Exit(1)
for root_file in root_files:
convert_steps_to_parquet(
root_file,
output_path=output,
batch_size=batch_size,
tree_name=tree,
compression=compression_value,
)
return
if output is not None:
typer.echo(
"error: --output cannot be combined with --jobs > 1 "
"(destinations are derived from --dataset-root/--schema)",
err=True,
)
raise typer.Exit(1)
run_parallel_job(
[str(f) for f in root_files],
jobs=jobs,
dataset_root=dataset_root,
schema=schema,
batch_size=batch_size,
tree=tree,
compression=compression_value,
)
@app.command()
def migrate(
root: Annotated[
Path, typer.Argument(help="Dataset root to migrate in place")
] = _DATASET_ROOT_DEFAULT,
execute: Annotated[
bool,
typer.Option(
"--execute",
help="Actually move/copy files and write manifests (default: dry run)",
),
] = False,
copy: Annotated[
bool,
typer.Option(
"--copy",
help="Copy instead of move, leaving the originals in place "
"(e.g. if another process is still reading them)",
),
] = False,
) -> None:
"""One-time migration into the versioned raw/processed/pools/derived layout."""
run_migration(str(root), execute=execute, copy=copy)
@app.command("bump-gen")
def bump_gen(
reason: Annotated[str, typer.Option("--reason", help="Why this gen exists")],
kind: Annotated[
str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
] = "steps",
by: Annotated[
Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
] = None,
date: Annotated[
Optional[str],
typer.Option("--date", help="Override date (default: today, ISO)"),
] = None,
to: Annotated[
Optional[str],
typer.Option(
"--to",
metavar="genN",
help="Target gen tag (default: one past the current highest)",
),
] = None,
execute: Annotated[
bool, typer.Option("--execute", help="Apply (default: dry run)")
] = False,
root: Annotated[
Path, typer.Option("--root", help="Dataset root")
] = _DATASET_ROOT_DEFAULT,
) -> None:
"""Cut a new raw generation."""
run_bump_gen(
kind=kind,
reason=reason,
by=by,
date=date,
execute=execute,
root=str(root),
to=to,
)
@app.command("bump-schema")
def bump_schema(
gen: Annotated[str, typer.Option("--gen", help="Existing gen tag, e.g. gen1")],
reason: Annotated[str, typer.Option("--reason", help="Why this schema exists")],
kind: Annotated[
str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
] = "steps",
by: Annotated[
Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
] = None,
date: Annotated[
Optional[str],
typer.Option("--date", help="Override date (default: today, ISO)"),
] = None,
to: Annotated[
Optional[str],
typer.Option(
"--to",
metavar="schemaN",
help="Target schema tag (default: one past the current highest)",
),
] = None,
execute: Annotated[
bool, typer.Option("--execute", help="Apply (default: dry run)")
] = False,
root: Annotated[
Path, typer.Option("--root", help="Dataset root")
] = _DATASET_ROOT_DEFAULT,
) -> None:
"""Cut a new schema within a gen."""
run_bump_schema(
kind=kind,
gen=gen,
reason=reason,
by=by,
date=date,
execute=execute,
root=str(root),
to=to,
)
@app.command()
def status(
root: Annotated[
Path, typer.Option("--root", help="Dataset root")
] = _DATASET_ROOT_DEFAULT,
) -> None:
"""List existing gens/schemas per kind."""
run_status(str(root))
@app.command("update-manifest")
def update_manifest(
manifests: Annotated[
list[Path], typer.Argument(help="One or more .manifest files to update")
],
schema: Annotated[
Optional[str],
typer.Option(
"--schema",
metavar="schemaN",
help="Target schema tag (default: highest schema found in the same gen dir)",
),
] = None,
gen: Annotated[
Optional[str],
typer.Option(
"--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"
),
] = None,
execute: Annotated[
bool,
typer.Option("--execute", help="Write updated manifests (default: dry run)"),
] = False,
) -> None:
"""Repoint manifest(s) to a new gen and/or schema, verifying all target files exist."""
run_update_manifest(
[str(m) for m in manifests], schema=schema, execute=execute, gen=gen
)
@app.command("create-manifest")
def create_manifest(
files: Annotated[list[Path], typer.Argument(help="Parquet files to include")],
output: Annotated[
Optional[Path],
typer.Option("--output", "-o", help="Explicit path for the new .manifest file"),
] = None,
pool: Annotated[
Optional[str],
typer.Option(
"--pool",
metavar="DETECTOR",
help="Detector name; combined with --type and --root to form "
"<root>/pools/<detector>/<type>.manifest",
),
] = None,
type_: Annotated[
Optional[PoolType],
typer.Option(
"--type", help="Pool type — full, holdout, or dev (required with --pool)"
),
] = None,
root: Annotated[
Path, typer.Option("--root", help="Dataset root (used with --pool)")
] = _DATASET_ROOT_DEFAULT,
execute: Annotated[
bool, typer.Option("--execute", help="Write the manifest (default: dry run)")
] = False,
) -> None:
"""Create a new manifest from a list of parquet files."""
run_create_manifest(
[str(f) for f in files],
execute=execute,
output=str(output) if output is not None else None,
pool=pool,
type_=type_.value if type_ is not None else None,
root=str(root),
)
@app.command("make-root")
def make_root(
executable: Annotated[
Path, typer.Option("--executable", help="Built minicalosim run_* executable")
],
detector: Annotated[
list[str],
typer.Option(
"--detector",
metavar="NAME[:CONFIG]",
help="Dataset detector label, optionally with ':CONFIG' to pass as the "
"executable's config-name argument (e.g. sampling_pb_scint:pb_scint). "
"Omit ':CONFIG' for executables that take no config selector (e.g. run_pbwo4). "
"Repeatable.",
),
],
num_files: Annotated[
int, typer.Option("--num-files", help="New shards to create per detector")
],
events_per_file: Annotated[
int, typer.Option("--events-per-file", help="nEvents passed to the executable")
],
gen: Annotated[
str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")
],
kind: Annotated[
str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
] = "steps",
dataset_root: Annotated[
Path, typer.Option("--dataset-root", help="Dataset root")
] = _DATASET_ROOT_DEFAULT,
jobs: Annotated[
int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")
] = 4,
execute: Annotated[
bool,
typer.Option(
"--execute", help="Actually run jobs (default: dry run / print plan)"
),
] = False,
) -> None:
"""Generate new ROOT shards via a minicalosim executable."""
run_make_root(
executable=executable,
detector=detector,
num_files=num_files,
events_per_file=events_per_file,
kind=kind,
gen=gen,
dataset_root=str(dataset_root),
jobs=jobs,
execute=execute,
)
@app.command("hparam-scan")
def hparam_scan(
data: Annotated[str, typer.Option("--data")] = DATA_DEFAULT,
scan_dir: Annotated[str, typer.Option("--scan-dir")] = SCAN_DIR_DEFAULT,
seed: Annotated[int, typer.Option("--seed")] = 0,
dry_run: Annotated[bool, typer.Option("--dry-run")] = False,
) -> None:
"""Grid-scan dropout x n_blocks x hidden_dim via sequential `giant train` runs."""
run_hparam_scan(data=data, scan_dir=scan_dir, seed=seed, dry_run=dry_run)
if __name__ == "__main__":
app()
scripts/hparam_scan.py
+14 -25
View File
@@ -1,4 +1,3 @@
#!/usr/bin/env python3
"""Hyperparameter scan over dropout x n_blocks x hidden_dim.
Runs `giant train` sequentially (this machine has a single GPU) for every
@@ -6,13 +5,9 @@ combination, plus one extra run at the default architecture with a higher
learning rate. Runs are shuffled so the parameter space gets coarse coverage
early rather than exhausting one corner of the grid first.
Usage:
uv run python scripts/hparam_scan.py
uv run python scripts/hparam_scan.py --dry-run
uv run python scripts/hparam_scan.py --seed 1 --data /path/to/parquet
See `uv run dwarf hparam-scan --help` for the CLI.
"""
import argparse
import csv
import itertools
import os
@@ -87,31 +82,29 @@ def append_summary(summary_path: Path, row: dict) -> None:
writer.writerow(row)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--data", default=DATA_DEFAULT)
parser.add_argument("--scan-dir", default=SCAN_DIR_DEFAULT)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--dry-run", action="store_true")
args = parser.parse_args()
def run_hparam_scan(
data: str = DATA_DEFAULT,
scan_dir: str = SCAN_DIR_DEFAULT,
seed: int = 0,
dry_run: bool = False,
) -> None:
runs = build_runs(seed)
scan_dir_path = Path(scan_dir)
runs = build_runs(args.seed)
scan_dir = Path(args.scan_dir)
if args.dry_run:
if dry_run:
for i, run in enumerate(runs, 1):
print(f"[{i}/{len(runs)}] {run_name(run)}")
return
scan_dir.mkdir(parents=True, exist_ok=True)
summary_path = scan_dir / "scan_summary.csv"
scan_dir_path.mkdir(parents=True, exist_ok=True)
summary_path = scan_dir_path / "scan_summary.csv"
env = os.environ.copy()
env["TQDM_DISABLE"] = "1"
for i, run in enumerate(runs, 1):
name = run_name(run)
out_dir = scan_dir / name
out_dir = scan_dir_path / name
metrics_path = out_dir / "metrics.csv"
last_ckpt = out_dir / "last.pt"
@@ -124,7 +117,7 @@ def main() -> None:
cmd = [
"giant",
"train",
args.data,
data,
"--mode",
"flow",
"--epochs",
@@ -186,7 +179,3 @@ def main() -> None:
print(
f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log"
)
if __name__ == "__main__":
main()
scripts/migrate_geant_steps.py
+14 -39
View File
@@ -1,4 +1,3 @@
#!/usr/bin/env python3
"""One-time migration of /ceph/lbogner/geant_steps into the versioned layout:
raw/<kind>/<gen>/<detector>/shard-NNN.root
@@ -13,12 +12,13 @@ up pool membership from POOL_ASSIGNMENT — decoupling "where it sits today" fro
"which pool it belongs to".
Defaults to a dry run (prints the planned moves and manifest contents). Pass
--execute to actually move files and write manifests. Pass --copy as well to
copy instead of move, leaving the original files in place — e.g. if another
process is still reading them from their current location.
execute=True to actually move files and write manifests, and copy=True to copy
instead of move, leaving the original files in place — e.g. if another process
is still reading them from their current location.
See `uv run dwarf migrate --help` for the CLI.
"""
import argparse
import os
import re
import shutil
@@ -181,39 +181,18 @@ def plan_manifests(src_root: Path) -> dict[Path, list[str]]:
}
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"root",
nargs="?",
default="/ceph/lbogner/geant_steps",
help="Dataset root to migrate in place (default: /ceph/lbogner/geant_steps)",
)
parser.add_argument(
"--execute",
action="store_true",
help="Actually move files and write manifests (default: dry run / print plan only)",
)
parser.add_argument(
"--copy",
action="store_true",
help="Copy instead of move, leaving the originals in place "
"(e.g. if another process is still reading them). Implies the legacy "
"train/ etc. directories are left as-is too, since they won't be empty.",
)
args = parser.parse_args()
src_root = Path(args.root)
def run_migration(root: str, execute: bool, copy: bool) -> None:
src_root = Path(root)
if not src_root.is_dir():
parser.error(f"{src_root} is not a directory")
raise SystemExit(f"error: {src_root} is not a directory")
moves, unrecognized = plan_moves(src_root)
manifests = plan_manifests(src_root)
verb = "COPY" if args.copy else "MOVE"
print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ({verb}): {src_root} ===\n")
verb = "COPY" if copy else "MOVE"
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ({verb}): {src_root} ===\n")
print(f"-- {len(moves)} file {'copy' if args.copy else 'move'}(s) --")
print(f"-- {len(moves)} file {'copy' if copy else 'move'}(s) --")
for src, dst in moves:
print(f" {src.relative_to(src_root)} -> {dst.relative_to(src_root)}")
@@ -228,11 +207,11 @@ def main() -> None:
for path in unrecognized:
print(f" {path.relative_to(src_root)}")
if not args.execute:
if not execute:
print("\nDry run only — pass --execute to apply.")
return
transfer = shutil.copy2 if args.copy else shutil.move
transfer = shutil.copy2 if copy else shutil.move
for src, dst in moves:
if dst.exists():
raise FileExistsError(f"refusing to overwrite existing file: {dst}")
@@ -257,7 +236,7 @@ def main() -> None:
# empty since their contents were classified by filename, not location —
# remove them, but only if a move actually emptied them. In --copy mode the
# originals are still there by design, so leave these alone entirely.
if not args.copy:
if not copy:
for stale_dir in ("train", "sampling_train/small", "sampling_train"):
d = src_root / stale_dir
if d.is_dir() and not any(d.iterdir()):
@@ -265,7 +244,3 @@ def main() -> None:
print(f"removed now-empty directory: {d.relative_to(src_root)}")
print("\nDone.")
if __name__ == "__main__":
main()
scripts/steps_to_parquet.py
+1 -54
View File
@@ -1,14 +1,8 @@
#!/usr/bin/env python3
"""Convert the Steps tree from a ROOT file to Parquet.
Usage:
uv run python steps_to_parquet.py input.root
uv run python steps_to_parquet.py input.root -o output.parquet
uv run python steps_to_parquet.py input.root --batch-size "200 MB" --tree Hits
uv run python steps_to_parquet.py input1.root input2.root input3.root
See `uv run dwarf convert --help` for the CLI.
"""
import argparse
from pathlib import Path
from typing import Literal
@@ -158,50 +152,3 @@ def convert_steps_to_parquet(
df.write_parquet(output_path, compression=compression)
print(f"done ({output_path.stat().st_size / 1e6:.1f} MB)")
return output_path
def main() -> None:
parser = argparse.ArgumentParser(
description="Convert a Steps (or any flat+jagged) tree in a ROOT file to Parquet."
)
parser.add_argument("root_files", nargs="+", help="Input ROOT file(s)")
parser.add_argument(
"-o",
"--output",
help="Output Parquet file (default: <input>.parquet). "
"Only valid with a single input file.",
)
parser.add_argument(
"--batch-size",
default="100 MB",
help="Uproot read batch size (default: '100 MB'). E.g. '50 MB', '500000' (rows).",
)
parser.add_argument(
"--tree",
default="Steps",
help="Tree name inside the ROOT file (default: Steps)",
)
parser.add_argument(
"--compression",
default="snappy",
choices=["snappy", "lz4", "zstd", "gzip", "none"],
help="Parquet compression codec (default: snappy)",
)
args = parser.parse_args()
if args.output is not None and len(args.root_files) > 1:
parser.error("--output can only be used with a single input file")
compression = "uncompressed" if args.compression == "none" else args.compression
for root_file in args.root_files:
convert_steps_to_parquet(
root_file,
output_path=args.output,
batch_size=args.batch_size,
tree_name=args.tree,
compression=compression,
)
if __name__ == "__main__":
main()
scripts/steps_to_parquet_parallel.py
+50 -88
View File
@@ -1,34 +1,27 @@
#!/usr/bin/env python3
"""Convert many ROOT files to Parquet by fanning out to steps_to_parquet.py.
"""Convert many ROOT files to Parquet by fanning out to `dwarf convert`.
steps_to_parquet.py itself converts a list of files one at a time; this wraps
it to run up to --jobs conversions concurrently, each as its own subprocess
(invoked with the same Python executable running this script, so it picks up
A single `dwarf convert` call converts a list of files one at a time; this
module runs up to --jobs conversions concurrently, each as its own `dwarf
convert` subprocess (invoked via `python -m scripts.dwarf`, so it picks up
the active venv/uv environment automatically).
Inputs must live under <dataset-root>/raw/<kind>/<gen>/<detector>/<file>.root
(see scripts/migrate_geant_steps.py) — each is written to the matching
processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet, where <schema>
defaults to the highest schemaN already under processed/<kind>/<gen>/ (pass
--schema to pick a specific one, e.g. one just created by
bump_dataset_version.py bump-schema). A file that doesn't fit that layout is
rejected up front, before any conversion runs.
--schema to pick a specific one, e.g. one just created by `dwarf bump-schema`).
A file that doesn't fit that layout is rejected up front, before any
conversion runs.
Usage:
steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/shard-000.root raw/steps/gen1/pbwo4/shard-001.root
steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/*.root --jobs 8
steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/*.root --schema schema2
See `uv run dwarf convert --help` for the CLI.
"""
import argparse
import re
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
_STEPS_TO_PARQUET = Path(__file__).resolve().parent / "steps_to_parquet.py"
# Must match scripts/bump_dataset_version.py's GEN_RE / SCHEMA_RE.
GEN_RE = re.compile(r"^gen\d+$")
SCHEMA_RE = re.compile(r"^schema(\d+)$")
@@ -90,17 +83,19 @@ def resolve_destination(
return processed_gen_dir / schema_tag / detector / f"{shard_stem}.parquet"
_DWARF_CONVERT_CMD = [sys.executable, "-m", "scripts.dwarf", "convert"]
def _convert_one(
root_file: str,
batch_size: str,
tree: str,
compression: str,
steps_to_parquet_path: Path,
output_path: Path | None,
cmd_prefix: list[str],
) -> tuple[str, int, str, str]:
cmd = [
sys.executable,
str(steps_to_parquet_path),
*cmd_prefix,
root_file,
"--batch-size",
batch_size,
@@ -122,20 +117,25 @@ def run_parallel(
batch_size: str = "100 MB",
tree: str = "Steps",
compression: str = "snappy",
steps_to_parquet_path: Path = _STEPS_TO_PARQUET,
output_for: dict[str, Path] | None = None,
cmd_prefix: list[str] | None = None,
) -> list[tuple[str, int, str, str]]:
"""Run one steps_to_parquet.py subprocess per file, up to *jobs* at a time.
"""Run one `dwarf convert` subprocess per file, up to *jobs* at a time.
*output_for*, if given, maps each root_file to the parquet path it should
be written to (passed through as steps_to_parquet.py's --output); files
missing from the map fall back to steps_to_parquet.py's own default
(parquet written next to the input .root).
be written to (passed through as `dwarf convert`'s --output); files
missing from the map fall back to `dwarf convert`'s own default (parquet
written next to the input .root).
*cmd_prefix* overrides the subprocess command run per file (defaults to
`python -m scripts.dwarf convert`) — used by tests to substitute a fake
conversion script.
Returns one (root_file, returncode, stdout, stderr) tuple per file, in
completion order (not necessarily input order).
"""
output_for = output_for or {}
cmd_prefix = cmd_prefix if cmd_prefix is not None else _DWARF_CONVERT_CMD
results = []
with ThreadPoolExecutor(max_workers=jobs) as pool:
futures = {
@@ -145,8 +145,8 @@ def run_parallel(
batch_size,
tree,
compression,
steps_to_parquet_path,
output_for.get(f),
cmd_prefix,
): f
for f in root_files
}
@@ -162,90 +162,52 @@ def run_parallel(
return results
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description="Convert many ROOT files to Parquet in parallel via steps_to_parquet.py."
)
parser.add_argument("root_files", nargs="+", help="Input ROOT file(s)")
parser.add_argument(
"-j",
"--jobs",
type=int,
default=4,
help="Number of conversions to run in parallel (default: 4)",
)
parser.add_argument(
"--batch-size",
default="100 MB",
help="Uproot read batch size (default: '100 MB'). E.g. '50 MB', '500000' (rows).",
)
parser.add_argument(
"--tree",
default="Steps",
help="Tree name inside the ROOT file (default: Steps)",
)
parser.add_argument(
"--compression",
default="snappy",
choices=["snappy", "lz4", "zstd", "gzip", "none"],
help="Parquet compression codec (default: snappy)",
)
parser.add_argument(
"--dataset-root",
default="/ceph/lbogner/geant_steps",
help="Dataset root containing raw/ and processed/ (default: /ceph/lbogner/geant_steps)",
)
parser.add_argument(
"--schema",
default=None,
help="Schema tag to write parquets under, e.g. schema2 "
"(default: highest schemaN already under processed/<kind>/<gen>/)",
)
return parser
def main() -> None:
parser = build_parser()
args = parser.parse_args()
if args.jobs < 1:
parser.error("--jobs must be >= 1")
dataset_root = Path(args.dataset_root)
def run_parallel_job(
root_files: list[str],
jobs: int,
dataset_root: Path,
schema: str | None,
batch_size: str,
tree: str,
compression: str,
) -> None:
"""Resolve each file's dataset-layout destination, convert in parallel, and
report results. Exits the process (via SystemExit) on destination or
conversion failure — this is the top-level entry point `dwarf convert`
delegates to when --jobs > 1."""
output_for: dict[str, Path] = {}
errors: list[str] = []
for f in args.root_files:
for f in root_files:
try:
output_for[f] = resolve_destination(Path(f), dataset_root, args.schema)
output_for[f] = resolve_destination(Path(f), dataset_root, schema)
except DestinationError as exc:
errors.append(str(exc))
if errors:
for err in errors:
print(f"error: {err}", file=sys.stderr)
sys.exit(1)
raise SystemExit(1)
for root_file, dest in output_for.items():
print(f"{root_file} -> {dest}")
results = run_parallel(
args.root_files,
jobs=args.jobs,
batch_size=args.batch_size,
tree=args.tree,
compression=args.compression,
root_files,
jobs=jobs,
batch_size=batch_size,
tree=tree,
compression=compression,
output_for=output_for,
)
failures = [root_file for root_file, code, _, _ in results if code != 0]
if failures:
print(f"\n{len(failures)} of {len(results)} conversion(s) failed:", file=sys.stderr)
print(
f"\n{len(failures)} of {len(results)} conversion(s) failed:",
file=sys.stderr,
)
for root_file in failures:
print(f" {root_file}", file=sys.stderr)
sys.exit(1)
raise SystemExit(1)
print(f"\nAll {len(results)} conversion(s) completed.")
if __name__ == "__main__":
main()