Merge energy-conservation-poc into phase2-secondary-prediction
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:
@@ -11,6 +11,9 @@ uv sync --extra cpu --extra dev # add dev extras (pytest, etc.)
|
||||
pytest # run tests
|
||||
giant train path/to/steps.parquet --mode flow # train (flow matching)
|
||||
giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline)
|
||||
dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen,
|
||||
# bump-schema, status, update-manifest, create-manifest,
|
||||
# make-root, hparam-scan (see scripts/dwarf.py)
|
||||
```
|
||||
|
||||
`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x; newer torch requires newer NVIDIA drivers). Plain `uv sync` with no extra will not install torch at all; uv has no concept of a "default extra", so `--extra cpu` should always be included unless you need GPU support.
|
||||
|
||||
@@ -34,7 +34,7 @@ The model is developed in two phases:
|
||||
|
||||
## Data
|
||||
|
||||
Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `uv run steps-to-parquet`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle) to avoid leaking correlated steps from the same shower.
|
||||
Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `uv run dwarf convert`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle) to avoid leaking correlated steps from the same shower.
|
||||
|
||||
## Project structure
|
||||
|
||||
@@ -55,13 +55,16 @@ giant/
|
||||
│ ├── validate.py # step-level marginal + KL-divergence validation
|
||||
│ ├── analysis.py # notebook diagnostics: marginals, correlations, constraint checks
|
||||
│ └── cli.py # `giant train` / `giant predict` Typer app
|
||||
├── scripts/ # also exposed as uv entry points, e.g. `uv run steps-to-parquet`
|
||||
│ ├── steps_to_parquet.py # ROOT → parquet conversion (uproot/awkward/polars)
|
||||
│ ├── steps_to_parquet_parallel.py # fan out steps_to_parquet.py over several ROOT files
|
||||
│ ├── migrate_geant_steps.py # one-time move into the raw/processed/pools/derived layout
|
||||
│ ├── bump_dataset_version.py # cut a new raw gen or parquet schema, with a logged reason
|
||||
│ ├── create_root_files.py # generate new ROOT shards via a minicalosim executable
|
||||
│ └── hparam_scan.py # hyperparameter grid scan over `giant train` runs
|
||||
├── scripts/ # dataset/tooling logic, unified under the `dwarf` CLI (`uv run dwarf --help`)
|
||||
│ ├── dwarf.py # Typer app: convert, migrate, bump-gen, bump-schema, status,
|
||||
│ │ # update-manifest, create-manifest, make-root, hparam-scan
|
||||
│ ├── steps_to_parquet.py # ROOT → parquet conversion (uproot/awkward/polars) — `dwarf convert`
|
||||
│ ├── steps_to_parquet_parallel.py # fan out conversion over several ROOT files — `dwarf convert --jobs N`
|
||||
│ ├── migrate_geant_steps.py # one-time move into the raw/processed/pools/derived layout — `dwarf migrate`
|
||||
│ ├── bump_dataset_version.py # cut a new raw gen or parquet schema, with a logged reason —
|
||||
│ │ # `dwarf bump-gen` / `bump-schema` / `status` / `update-manifest` / `create-manifest`
|
||||
│ ├── create_root_files.py # generate new ROOT shards via a minicalosim executable — `dwarf make-root`
|
||||
│ └── hparam_scan.py # hyperparameter grid scan over `giant train` runs — `dwarf hparam-scan`
|
||||
└── tests/
|
||||
```
|
||||
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
"""Compare energy-conservation PoC at 10 vs 20 ODE steps.
|
||||
|
||||
Runs the same event-level energy-budget analysis as
|
||||
`analysis/export_energy_conservation_poc.py` on both predict outputs and prints a
|
||||
side-by-side table. Also regenerates the two incident-energy comparison plots for
|
||||
the 20-step run (prefix `giant-energy-conservation-poc-ode20-`).
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from giant.analysis import _edep_pl, _hist_edges
|
||||
|
||||
FILES = {
|
||||
"10-step (baseline)": "/home/lars/Programming/giant/9879e806-5e88-4b06-b1fa-0e61de9cda6f.parquet",
|
||||
"20-step": "/home/lars/Programming/giant/0e236919-65ae-4b9b-9957-d31a8211aca4.parquet",
|
||||
}
|
||||
OUT = Path("/home/lars/knowledge-base/meta/attachments")
|
||||
PREFIX20 = "giant-energy-conservation-poc-ode20"
|
||||
|
||||
|
||||
def per_event(file: str) -> dict:
|
||||
pe = (
|
||||
pl.scan_parquet(file)
|
||||
.group_by("event_id")
|
||||
.agg(
|
||||
pl.col("pre_E").max().alias("primary_E"),
|
||||
_edep_pl("true").sum().alias("real_total_edep"),
|
||||
_edep_pl("pred").sum().alias("gen_total_edep"),
|
||||
)
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
primary_E = pe["primary_E"].to_numpy()
|
||||
assert np.unique(primary_E).size == 1, "expected a single fixed incident energy"
|
||||
E0 = float(primary_E[0])
|
||||
real = pe["real_total_edep"].to_numpy()
|
||||
gen = pe["gen_total_edep"].to_numpy()
|
||||
n_steps = (
|
||||
pl.scan_parquet(file).select(pl.len()).collect(engine="streaming").item()
|
||||
)
|
||||
return {
|
||||
"E0": E0,
|
||||
"n_events": pe.height,
|
||||
"n_rows": n_steps,
|
||||
"real": real,
|
||||
"gen": gen,
|
||||
}
|
||||
|
||||
|
||||
results = {name: per_event(f) for name, f in FILES.items()}
|
||||
|
||||
|
||||
def fmt_row(label, fn):
|
||||
cells = " ".join(f"{fn(r):>14}" for r in results.values())
|
||||
print(f"{label:<32}{cells}")
|
||||
|
||||
|
||||
print("=" * 80)
|
||||
header = " ".join(f"{name:>14}" for name in results)
|
||||
print(f"{'metric':<32}{header}")
|
||||
print("-" * 80)
|
||||
fmt_row("n_events", lambda r: r["n_events"])
|
||||
fmt_row("n_rows (steps)", lambda r: r["n_rows"])
|
||||
fmt_row("E0 [MeV]", lambda r: f"{r['E0']:.1f}")
|
||||
print("-- REAL --")
|
||||
fmt_row("real mean [MeV]", lambda r: f"{r['real'].mean():.3f}")
|
||||
fmt_row("real std [MeV]", lambda r: f"{r['real'].std():.3f}")
|
||||
fmt_row("real sigma/mu", lambda r: f"{r['real'].std() / r['real'].mean():.4f}")
|
||||
print("-- GENERATED --")
|
||||
fmt_row("gen mean [MeV]", lambda r: f"{r['gen'].mean():.3f}")
|
||||
fmt_row("gen std [MeV]", lambda r: f"{r['gen'].std():.3f}")
|
||||
fmt_row("gen sigma/mu", lambda r: f"{r['gen'].std() / r['gen'].mean():.4f}")
|
||||
fmt_row("gen max [MeV]", lambda r: f"{r['gen'].max():.3f}")
|
||||
fmt_row("gen mean/E0", lambda r: f"{r['gen'].mean() / r['E0']:.4f}")
|
||||
fmt_row("gen max/E0", lambda r: f"{r['gen'].max() / r['E0']:.4f}")
|
||||
fmt_row("gen p99/E0", lambda r: f"{np.quantile(r['gen'], 0.99) / r['E0']:.4f}")
|
||||
fmt_row("frac events gen>E0", lambda r: f"{np.mean(r['gen'] > r['E0']):.4f}")
|
||||
def disp_ratio(r):
|
||||
return (r["gen"].std() / r["gen"].mean()) / (r["real"].std() / r["real"].mean())
|
||||
|
||||
|
||||
fmt_row("dispersion ratio gen/real", lambda r: f"{disp_ratio(r):.2f}x")
|
||||
print("=" * 80)
|
||||
|
||||
# --- plots for the 20-step run (mirror the baseline export) ---
|
||||
r = results["20-step"]
|
||||
E0, real_tot, gen_tot = r["E0"], r["real"], r["gen"]
|
||||
|
||||
fig, ax = plt.subplots(figsize=(6, 4))
|
||||
edges = _hist_edges(real_tot, gen_tot, bins=50).tolist()
|
||||
ax.hist(
|
||||
real_tot, bins=edges, density=True, histtype="step",
|
||||
label=f"real (σ/μ={real_tot.std() / real_tot.mean():.3f})",
|
||||
)
|
||||
ax.hist(
|
||||
gen_tot, bins=edges, density=True, histtype="step",
|
||||
label=f"generated (σ/μ={gen_tot.std() / gen_tot.mean():.3f}, "
|
||||
f"{np.mean(gen_tot > E0):.1%} > E0)",
|
||||
)
|
||||
ax.axvline(E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV")
|
||||
ax.set_yscale("log")
|
||||
ax.set_xlabel("total deposited energy per event [MeV]")
|
||||
ax.set_title("20 ODE steps")
|
||||
ax.legend(fontsize=8)
|
||||
fig.tight_layout()
|
||||
fig.savefig(OUT / f"{PREFIX20}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight")
|
||||
|
||||
fig, ax = plt.subplots(figsize=(6, 4))
|
||||
ratio_real = real_tot / E0
|
||||
ratio_gen = gen_tot / E0
|
||||
edges = _hist_edges(ratio_real, ratio_gen, bins=60).tolist()
|
||||
ax.hist(ratio_real, bins=edges, density=True, histtype="step", label="real")
|
||||
ax.hist(ratio_gen, bins=edges, density=True, histtype="step", label="generated")
|
||||
ax.axvline(1.0, color="k", linestyle="--", linewidth=1, label="conservation limit (=1)")
|
||||
ax.set_yscale("log")
|
||||
ax.set_xlabel("total deposited energy / incident energy, per event")
|
||||
ax.set_title("20 ODE steps")
|
||||
ax.legend(fontsize=8)
|
||||
fig.tight_layout()
|
||||
fig.savefig(OUT / f"{PREFIX20}-event-energy-ratio.png", dpi=150, bbox_inches="tight")
|
||||
|
||||
# --- overlay: generated total-edep, 10 vs 20 steps, against real+E0 ---
|
||||
fig, ax = plt.subplots(figsize=(6, 4))
|
||||
all_arrays = [results["10-step (baseline)"]["real"]] + [
|
||||
r2["gen"] for r2 in results.values()
|
||||
]
|
||||
edges = _hist_edges(*all_arrays, bins=60).tolist()
|
||||
ax.hist(
|
||||
results["20-step"]["real"], bins=edges, density=True, histtype="step",
|
||||
color="k", label="real",
|
||||
)
|
||||
for name, r2 in results.items():
|
||||
ax.hist(
|
||||
r2["gen"], bins=edges, density=True, histtype="step",
|
||||
label=f"gen {name} ({np.mean(r2['gen'] > r2['E0']):.1%} > E0)",
|
||||
)
|
||||
ax.axvline(E0, color="gray", linestyle="--", linewidth=1, label=f"E0={E0:.0f} MeV")
|
||||
ax.set_yscale("log")
|
||||
ax.set_xlabel("total deposited energy per event [MeV]")
|
||||
ax.set_title("Generated event energy: 10 vs 20 ODE steps")
|
||||
ax.legend(fontsize=8)
|
||||
fig.tight_layout()
|
||||
fig.savefig(OUT / f"{PREFIX20}-compare-event-total-energy.png", dpi=150, bbox_inches="tight")
|
||||
|
||||
print("DONE")
|
||||
@@ -0,0 +1,71 @@
|
||||
"""Per-step marginal KL(real||gen) per target dim, 10 vs 20 ODE steps.
|
||||
|
||||
Streaming histogram over shared bin edges (computed from the true distribution),
|
||||
so the two runs are directly comparable dim-by-dim.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from giant.analysis import RAW_TARGET_NAMES, _kl_from_counts, _raw_dim_expr
|
||||
|
||||
FILES = {
|
||||
"10-step": "/home/lars/Programming/giant/9879e806-5e88-4b06-b1fa-0e61de9cda6f.parquet",
|
||||
"20-step": "/home/lars/Programming/giant/0e236919-65ae-4b9b-9957-d31a8211aca4.parquet",
|
||||
}
|
||||
BINS = 100
|
||||
|
||||
# Fixed shared edges from the true distribution (identical across both files), using
|
||||
# robust quantiles to avoid a few outliers dominating the range.
|
||||
base = FILES["10-step"]
|
||||
edges = {}
|
||||
for j, name in enumerate(RAW_TARGET_NAMES):
|
||||
lo, hi = (
|
||||
pl.scan_parquet(base)
|
||||
.select(
|
||||
_raw_dim_expr("true", j).quantile(0.001).alias("lo"),
|
||||
_raw_dim_expr("true", j).quantile(0.999).alias("hi"),
|
||||
)
|
||||
.collect(engine="streaming")
|
||||
.row(0)
|
||||
)
|
||||
if not (hi - lo > 1e-9):
|
||||
lo, hi = lo - 0.5, hi + 0.5
|
||||
edges[name] = np.linspace(lo, hi, BINS + 1)
|
||||
|
||||
|
||||
def counts(file, prefix, j, e):
|
||||
vals = (
|
||||
pl.scan_parquet(file)
|
||||
.select(_raw_dim_expr(prefix, j).alias("v"))
|
||||
.collect(engine="streaming")["v"]
|
||||
.to_numpy()
|
||||
)
|
||||
c, _ = np.histogram(vals, bins=e)
|
||||
return c
|
||||
|
||||
|
||||
kls = {name: {} for name in FILES}
|
||||
for j, name in enumerate(RAW_TARGET_NAMES):
|
||||
e = edges[name]
|
||||
real_c = counts(base, "true", j, e) # identical true dist across files
|
||||
for run, f in FILES.items():
|
||||
gen_c = counts(f, "pred", j, e)
|
||||
kls[run][name] = _kl_from_counts(real_c, gen_c)
|
||||
|
||||
print(f"{'dim':<14}{'KL 10-step':>14}{'KL 20-step':>14}{'ratio 20/10':>14}")
|
||||
print("-" * 56)
|
||||
tot = {"10-step": 0.0, "20-step": 0.0}
|
||||
for name in RAW_TARGET_NAMES:
|
||||
a, b = kls["10-step"][name], kls["20-step"][name]
|
||||
tot["10-step"] += a
|
||||
tot["20-step"] += b
|
||||
print(f"{name:<14}{a:>14.5f}{b:>14.5f}{b / a if a else float('nan'):>14.2f}")
|
||||
print("-" * 56)
|
||||
print(
|
||||
f"{'SUM':<14}{tot['10-step']:>14.5f}{tot['20-step']:>14.5f}"
|
||||
f"{tot['20-step'] / tot['10-step']:>14.2f}"
|
||||
)
|
||||
print(
|
||||
f"{'MEAN':<14}{tot['10-step'] / 9:>14.5f}{tot['20-step'] / 9:>14.5f}"
|
||||
)
|
||||
@@ -0,0 +1,101 @@
|
||||
"""One-off export of the energy-budget-violation plot for the
|
||||
energy-conservation PoC checkpoint (ALR simplex output space, retrained on
|
||||
the patched-Geant4 regenerated dataset). Not part of the package; run
|
||||
manually.
|
||||
|
||||
Unlike `plot_total_energy` (per-step conservation only, no reference to the
|
||||
fixed primary/incident energy), this adds an explicit comparison against
|
||||
`primary_E` (== max(pre_E) per event, since this PoC dataset uses a single
|
||||
fixed incident energy) to show event-level conservation violations that the
|
||||
per-step ALR simplex constraint does not prevent.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from giant.analysis import _edep_pl, _hist_edges
|
||||
|
||||
FILE = "/home/lars/Programming/giant/9879e806-5e88-4b06-b1fa-0e61de9cda6f.parquet"
|
||||
OUT = Path("/home/lars/knowledge-base/meta/attachments")
|
||||
PREFIX = "giant-energy-conservation-poc"
|
||||
|
||||
print("=== per-event primary energy + total edep ===")
|
||||
per_event = (
|
||||
pl.scan_parquet(FILE)
|
||||
.group_by("event_id")
|
||||
.agg(
|
||||
pl.col("pre_E").max().alias("primary_E"),
|
||||
_edep_pl("true").sum().alias("real_total_edep"),
|
||||
_edep_pl("pred").sum().alias("gen_total_edep"),
|
||||
)
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
print("n events:", per_event.height)
|
||||
|
||||
primary_E = per_event["primary_E"].to_numpy()
|
||||
real_tot = per_event["real_total_edep"].to_numpy()
|
||||
gen_tot = per_event["gen_total_edep"].to_numpy()
|
||||
|
||||
assert np.unique(primary_E).size == 1, "expected a single fixed incident energy"
|
||||
E0 = float(primary_E[0])
|
||||
print(f"fixed incident energy E0 = {E0} MeV")
|
||||
|
||||
print()
|
||||
print(
|
||||
f"real: mean={real_tot.mean():.3f} std={real_tot.std():.3f} "
|
||||
f"sigma/mu={real_tot.std() / real_tot.mean():.4f} max={real_tot.max():.3f} "
|
||||
f"frac>E0={np.mean(real_tot > E0):.4f}"
|
||||
)
|
||||
print(
|
||||
f"gen: mean={gen_tot.mean():.3f} std={gen_tot.std():.3f} "
|
||||
f"sigma/mu={gen_tot.std() / gen_tot.mean():.4f} max={gen_tot.max():.3f} "
|
||||
f"frac>E0={np.mean(gen_tot > E0):.4f}"
|
||||
)
|
||||
print(
|
||||
f"gen/E0 ratio: mean={np.mean(gen_tot / E0):.4f} max={np.max(gen_tot / E0):.4f} "
|
||||
f"p99={np.quantile(gen_tot / E0, 0.99):.4f}"
|
||||
)
|
||||
|
||||
print("=== plot: total edep per event, marked against incident energy ===")
|
||||
fig, ax = plt.subplots(figsize=(6, 4))
|
||||
edges = _hist_edges(real_tot, gen_tot, bins=50).tolist()
|
||||
ax.hist(
|
||||
real_tot,
|
||||
bins=edges,
|
||||
density=True,
|
||||
histtype="step",
|
||||
label=f"real (σ/μ={real_tot.std() / real_tot.mean():.3f})",
|
||||
)
|
||||
ax.hist(
|
||||
gen_tot,
|
||||
bins=edges,
|
||||
density=True,
|
||||
histtype="step",
|
||||
label=f"generated (σ/μ={gen_tot.std() / gen_tot.mean():.3f}, "
|
||||
f"{np.mean(gen_tot > E0):.1%} > E0)",
|
||||
)
|
||||
ax.axvline(E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV")
|
||||
ax.set_yscale("log")
|
||||
ax.set_xlabel("total deposited energy per event [MeV]")
|
||||
ax.legend(fontsize=8)
|
||||
fig.tight_layout()
|
||||
fig.savefig(OUT / f"{PREFIX}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight")
|
||||
|
||||
print("=== plot: total edep / incident energy ratio ===")
|
||||
fig, ax = plt.subplots(figsize=(6, 4))
|
||||
ratio_real = real_tot / E0
|
||||
ratio_gen = gen_tot / E0
|
||||
edges = _hist_edges(ratio_real, ratio_gen, bins=60).tolist()
|
||||
ax.hist(ratio_real, bins=edges, density=True, histtype="step", label="real")
|
||||
ax.hist(ratio_gen, bins=edges, density=True, histtype="step", label="generated")
|
||||
ax.axvline(1.0, color="k", linestyle="--", linewidth=1, label="conservation limit (=1)")
|
||||
ax.set_yscale("log")
|
||||
ax.set_xlabel("total deposited energy / incident energy, per event")
|
||||
ax.legend(fontsize=8)
|
||||
fig.tight_layout()
|
||||
fig.savefig(OUT / f"{PREFIX}-event-energy-ratio.png", dpi=150, bbox_inches="tight")
|
||||
|
||||
print("DONE")
|
||||
@@ -23,15 +23,23 @@ fig.savefig(OUT / f"{PREFIX}-event-total-length.png", dpi=150, bbox_inches="tigh
|
||||
|
||||
print("=== mean/median energy & length per step ===")
|
||||
fig = a.plot_mean_energy_per_step(obs)
|
||||
fig.savefig(OUT / f"{PREFIX}-event-mean-energy-per-step.png", dpi=150, bbox_inches="tight")
|
||||
fig.savefig(
|
||||
OUT / f"{PREFIX}-event-mean-energy-per-step.png", dpi=150, bbox_inches="tight"
|
||||
)
|
||||
fig = a.plot_mean_length_per_step(obs)
|
||||
fig.savefig(OUT / f"{PREFIX}-event-mean-length-per-step.png", dpi=150, bbox_inches="tight")
|
||||
fig.savefig(
|
||||
OUT / f"{PREFIX}-event-mean-length-per-step.png", dpi=150, bbox_inches="tight"
|
||||
)
|
||||
|
||||
print("=== longitudinal / transverse profiles ===")
|
||||
fig = a.plot_longitudinal_profile(obs)
|
||||
fig.savefig(OUT / f"{PREFIX}-event-longitudinal-profile.png", dpi=150, bbox_inches="tight")
|
||||
fig.savefig(
|
||||
OUT / f"{PREFIX}-event-longitudinal-profile.png", dpi=150, bbox_inches="tight"
|
||||
)
|
||||
fig = a.plot_transverse_profile(obs)
|
||||
fig.savefig(OUT / f"{PREFIX}-event-transverse-profile.png", dpi=150, bbox_inches="tight")
|
||||
fig.savefig(
|
||||
OUT / f"{PREFIX}-event-transverse-profile.png", dpi=150, bbox_inches="tight"
|
||||
)
|
||||
|
||||
print("=== shower-max depth ===")
|
||||
fig = a.plot_shower_max_depth(obs)
|
||||
@@ -45,8 +53,6 @@ fig = a.plot_pdg_length_share(pdg_table)
|
||||
fig.savefig(OUT / f"{PREFIX}-pdg-length-share.png", dpi=150, bbox_inches="tight")
|
||||
|
||||
print("=== summary stats ===")
|
||||
import numpy as np # noqa: E402
|
||||
|
||||
for label, real_col, gen_col in [
|
||||
("total_edep", "real_total_edep", "gen_total_edep"),
|
||||
("total_length", "real_total_length", "gen_total_length"),
|
||||
|
||||
@@ -17,8 +17,26 @@ edep_idx = a.RAW_TARGET_NAMES.index("edep")
|
||||
real = samples.real_raw[mask, edep_idx]
|
||||
gen = samples.gen_raw[mask, edep_idx]
|
||||
print("n photon rows:", mask.sum())
|
||||
print("real: mean", real.mean(), "std", real.std(), "max", real.max(), "frac==0", (real == 0).mean())
|
||||
print("gen: mean", gen.mean(), "std", gen.std(), "max", gen.max(), "frac==0", (gen == 0).mean())
|
||||
print(
|
||||
"real: mean",
|
||||
real.mean(),
|
||||
"std",
|
||||
real.std(),
|
||||
"max",
|
||||
real.max(),
|
||||
"frac==0",
|
||||
(real == 0).mean(),
|
||||
)
|
||||
print(
|
||||
"gen: mean",
|
||||
gen.mean(),
|
||||
"std",
|
||||
gen.std(),
|
||||
"max",
|
||||
gen.max(),
|
||||
"frac==0",
|
||||
(gen == 0).mean(),
|
||||
)
|
||||
for q in [0.5, 0.9, 0.99, 0.999]:
|
||||
print(f"q={q}: real={np.quantile(real, q):.4f} gen={np.quantile(gen, q):.4f}")
|
||||
|
||||
@@ -30,11 +48,29 @@ axes[0].set_yscale("log")
|
||||
axes[0].set_xlabel("edep (photons, pdg=22)")
|
||||
axes[0].legend()
|
||||
|
||||
axes[1].hist(real, bins=bins, alpha=0.6, label="real", density=True, cumulative=True, histtype="step")
|
||||
axes[1].hist(gen, bins=bins, alpha=0.6, label="gen", density=True, cumulative=True, histtype="step")
|
||||
axes[1].hist(
|
||||
real,
|
||||
bins=bins,
|
||||
alpha=0.6,
|
||||
label="real",
|
||||
density=True,
|
||||
cumulative=True,
|
||||
histtype="step",
|
||||
)
|
||||
axes[1].hist(
|
||||
gen,
|
||||
bins=bins,
|
||||
alpha=0.6,
|
||||
label="gen",
|
||||
density=True,
|
||||
cumulative=True,
|
||||
histtype="step",
|
||||
)
|
||||
axes[1].set_xlabel("edep (photons, pdg=22) - CDF")
|
||||
axes[1].legend()
|
||||
|
||||
fig.tight_layout()
|
||||
fig.savefig(OUT / "giant-h1024n8d0.1lr3e-4-photon-edep-zoom.png", dpi=150, bbox_inches="tight")
|
||||
fig.savefig(
|
||||
OUT / "giant-h1024n8d0.1lr3e-4-photon-edep-zoom.png", dpi=150, bbox_inches="tight"
|
||||
)
|
||||
print("saved photon-edep-zoom")
|
||||
|
||||
@@ -31,9 +31,16 @@ print("n photon rows:", df.height)
|
||||
|
||||
# Process order by abundance, so the legend is stable and the busiest on top.
|
||||
processes = (
|
||||
df.group_by("process").len().sort("len", descending=True).get_column("process").to_list()
|
||||
df.group_by("process")
|
||||
.len()
|
||||
.sort("len", descending=True)
|
||||
.get_column("process")
|
||||
.to_list()
|
||||
)
|
||||
series = {p: df.filter(pl.col("process") == p).get_column("edep_ev").to_numpy() for p in processes}
|
||||
series = {
|
||||
p: df.filter(pl.col("process") == p).get_column("edep_ev").to_numpy()
|
||||
for p in processes
|
||||
}
|
||||
|
||||
all_edep = df.get_column("edep_ev").to_numpy()
|
||||
lin_bins = np.linspace(0, np.quantile(all_edep, 0.999), 80)
|
||||
|
||||
@@ -36,7 +36,9 @@ ax.set_ylabel("density")
|
||||
ax.legend()
|
||||
|
||||
fig.tight_layout()
|
||||
fig.savefig(OUT / "giant-h1024n8d0.1lr3e-4-photon-edep-ev.png", dpi=150, bbox_inches="tight")
|
||||
fig.savefig(
|
||||
OUT / "giant-h1024n8d0.1lr3e-4-photon-edep-ev.png", dpi=150, bbox_inches="tight"
|
||||
)
|
||||
print("saved photon-edep-ev")
|
||||
|
||||
# Second version: log-spaced energy axis to expose the low-deposition structure.
|
||||
@@ -53,5 +55,9 @@ ax.set_ylabel("density")
|
||||
ax.legend()
|
||||
|
||||
fig.tight_layout()
|
||||
fig.savefig(OUT / "giant-h1024n8d0.1lr3e-4-photon-edep-ev-logx.png", dpi=150, bbox_inches="tight")
|
||||
fig.savefig(
|
||||
OUT / "giant-h1024n8d0.1lr3e-4-photon-edep-ev-logx.png",
|
||||
dpi=150,
|
||||
bbox_inches="tight",
|
||||
)
|
||||
print("saved photon-edep-ev-logx")
|
||||
|
||||
@@ -91,10 +91,22 @@ axes[0].set_yscale("log")
|
||||
axes[0].set_xlabel("edep (photons, pdg=22)")
|
||||
axes[0].legend()
|
||||
axes[1].hist(
|
||||
real, bins=bins, alpha=0.6, label="real", density=True, cumulative=True, histtype="step"
|
||||
real,
|
||||
bins=bins,
|
||||
alpha=0.6,
|
||||
label="real",
|
||||
density=True,
|
||||
cumulative=True,
|
||||
histtype="step",
|
||||
)
|
||||
axes[1].hist(
|
||||
gen, bins=bins, alpha=0.6, label="gen", density=True, cumulative=True, histtype="step"
|
||||
gen,
|
||||
bins=bins,
|
||||
alpha=0.6,
|
||||
label="gen",
|
||||
density=True,
|
||||
cumulative=True,
|
||||
histtype="step",
|
||||
)
|
||||
axes[1].set_xlabel("edep (photons, pdg=22) - CDF")
|
||||
axes[1].legend()
|
||||
|
||||
+28
-28
File diff suppressed because one or more lines are too long
+285
-287
@@ -1,15 +1,16 @@
|
||||
"""Notebook-friendly diagnostics for a trained model's sample quality.
|
||||
|
||||
Typical use from a Jupyter notebook::
|
||||
All checks here read from `giant predict --coord local` parquet output
|
||||
(`pred_*`/`true_*` columns, denormalized but still local-frame/log-scaled) —
|
||||
there is no on-the-fly (checkpoint + live sampler) path; generate predictions
|
||||
once via the CLI, then run every diagnostic below against that file::
|
||||
|
||||
from giant.analysis import load_model_bundle, make_val_loader, collect_samples
|
||||
from giant.analysis import load_predicted_local
|
||||
from giant.analysis import plot_marginals, plot_kl_bars, plot_correlation_matrices
|
||||
from giant.analysis import plot_pairwise, plot_direction_alignment
|
||||
from giant.analysis import plot_constraint_violations
|
||||
|
||||
bundle = load_model_bundle("runs/my_run/best.pt")
|
||||
val_loader = make_val_loader(bundle, "path/to/steps.parquet")
|
||||
samples = collect_samples(bundle, val_loader)
|
||||
samples = load_predicted_local("path/to/steps_predicted_local.parquet")
|
||||
|
||||
plot_marginals(samples, group_by="energy")
|
||||
plot_kl_bars(samples, group_by="energy")
|
||||
@@ -18,12 +19,6 @@ Typical use from a Jupyter notebook::
|
||||
plot_direction_alignment(samples)
|
||||
plot_constraint_violations(samples)
|
||||
|
||||
If predictions were already generated offline via `giant predict --coord local`,
|
||||
skip the checkpoint/model entirely and load the parquet directly::
|
||||
|
||||
from giant.analysis import load_predicted_local
|
||||
samples = load_predicted_local("path/to/steps_predicted_local.parquet")
|
||||
|
||||
(`--coord global` output isn't supported here — it has no ground-truth columns
|
||||
to compare against.)
|
||||
|
||||
@@ -77,11 +72,6 @@ Four tiers of checks, building on the aggregate marginal/KL check in
|
||||
surface covariate-shift failures that only appear under true rollout, only
|
||||
how well one-step generation reconstructs aggregate shower structure when
|
||||
fed real conditioning throughout.
|
||||
|
||||
`collect_samples` takes a `steps` argument (forwarded to the flow ODE
|
||||
integrator or, in ddim mode, the DDIM substep count) so a later
|
||||
sampler-step-count ablation can sweep it by calling this function repeatedly
|
||||
without any new plumbing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -93,29 +83,20 @@ import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import polars as pl
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
import pyarrow.parquet as pq
|
||||
|
||||
from giant.config import warn_if_checkpoint_config_mismatch
|
||||
from giant.constants import (
|
||||
LOCAL_TARGET_NAMES,
|
||||
PREDICT_COORD_METADATA_KEY,
|
||||
PREDICT_SCHEMA_VERSION,
|
||||
PREDICT_SCHEMA_VERSION_KEY,
|
||||
)
|
||||
from giant.data.dataset import make_event_split, StreamingStepsDataset
|
||||
from giant.data.loader import find_parquet_files, load_event_ids
|
||||
from giant.data.transforms import (
|
||||
Normalizer,
|
||||
energy_simplex_decode,
|
||||
inv_log_transform,
|
||||
reconstruct_post_pos,
|
||||
)
|
||||
from giant.model.network import DenoisingMLP
|
||||
from giant.model.schedule import CosineSchedule
|
||||
from giant.sample import sample_ddim, sample_ddpm, sample_flow
|
||||
from giant.validate import _histogram_kl
|
||||
|
||||
# The first 3 target dims are the scalar (non-direction) outputs. In raw/physical
|
||||
@@ -190,27 +171,6 @@ def _hist_edges(*arrays: np.ndarray, bins: int) -> np.ndarray:
|
||||
return np.linspace(lo, hi, bins + 1)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelBundle:
|
||||
model: torch.nn.Module
|
||||
cond_normalizer: Normalizer
|
||||
target_normalizer: Normalizer
|
||||
pdg_map: dict[int, int]
|
||||
mat_map: dict[str, int]
|
||||
model_config: dict
|
||||
mode: str
|
||||
schedule: CosineSchedule | None
|
||||
device: torch.device
|
||||
|
||||
@property
|
||||
def idx_to_pdg(self) -> dict[int, int]:
|
||||
return {v: k for k, v in self.pdg_map.items()}
|
||||
|
||||
@property
|
||||
def idx_to_mat(self) -> dict[int, str]:
|
||||
return {v: k for k, v in self.mat_map.items()}
|
||||
|
||||
|
||||
@dataclass
|
||||
class SampleCollection:
|
||||
cond_cont_raw: np.ndarray # (N, 9) denormalized conditioning (pre_E delogged)
|
||||
@@ -218,146 +178,6 @@ class SampleCollection:
|
||||
material: np.ndarray # (N,) raw material names
|
||||
real_raw: np.ndarray # (N, 9) denormalized + delogged real targets
|
||||
gen_raw: np.ndarray # (N, 9) denormalized + delogged generated targets
|
||||
# Normalized-space (model's native training space) targets — only available
|
||||
# when collected live with the normalizer (collect_samples). A predict
|
||||
# parquet has already been denormalized on disk with no normalizer
|
||||
# attached, so loaders built from one (e.g. load_predicted_local) leave
|
||||
# these as None rather than fabricate a value.
|
||||
real_norm: np.ndarray | None = None
|
||||
gen_norm: np.ndarray | None = None
|
||||
|
||||
|
||||
def load_model_bundle(
|
||||
ckpt_path: str | Path,
|
||||
mode: str = "flow",
|
||||
device: torch.device | None = None,
|
||||
) -> ModelBundle:
|
||||
"""Reconstruct a trained model and its normalizers/vocab from a training checkpoint."""
|
||||
warn_if_checkpoint_config_mismatch(ckpt_path)
|
||||
device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
|
||||
|
||||
pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
|
||||
mat_map = dict(ckpt["mat_map"])
|
||||
|
||||
model = DenoisingMLP(**ckpt["model_config"])
|
||||
model.load_state_dict(ckpt["model"])
|
||||
model.to(device).eval()
|
||||
|
||||
cond_normalizer = Normalizer.from_dict(ckpt["normalizer"]["cond"])
|
||||
target_normalizer = Normalizer.from_dict(ckpt["normalizer"]["target"])
|
||||
schedule = CosineSchedule().to(device) if mode != "flow" else None
|
||||
|
||||
return ModelBundle(
|
||||
model=model,
|
||||
cond_normalizer=cond_normalizer,
|
||||
target_normalizer=target_normalizer,
|
||||
pdg_map=pdg_map,
|
||||
mat_map=mat_map,
|
||||
model_config=ckpt["model_config"],
|
||||
mode=mode,
|
||||
schedule=schedule,
|
||||
device=device,
|
||||
)
|
||||
|
||||
|
||||
def make_val_loader(
|
||||
bundle: ModelBundle,
|
||||
data: str | Path,
|
||||
val_fraction: float = 0.1,
|
||||
seed: int = 42,
|
||||
batch_size: int = 4096,
|
||||
) -> DataLoader:
|
||||
"""Build a DataLoader over the val split, normalized with the bundle's fitted stats.
|
||||
|
||||
Loads the whole file into memory — fine for typical validation-set sizes; for
|
||||
very large datasets, build a StreamingStepsDataset directly (see giant.pipeline).
|
||||
"""
|
||||
files = find_parquet_files(data)
|
||||
all_event_ids = np.concatenate([load_event_ids(f) for f in files])
|
||||
_, val_events = make_event_split(all_event_ids, val_fraction=val_fraction, seed=seed)
|
||||
|
||||
val_ds = StreamingStepsDataset(
|
||||
files=files,
|
||||
split_events=val_events,
|
||||
pdg_map=bundle.pdg_map,
|
||||
mat_map=bundle.mat_map,
|
||||
cond_normalizer=bundle.cond_normalizer,
|
||||
target_normalizer=bundle.target_normalizer,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
)
|
||||
return DataLoader(val_ds, batch_size=None)
|
||||
|
||||
|
||||
def _to_raw_targets(
|
||||
target_norm: np.ndarray, normalizer: Normalizer, pre_E: np.ndarray
|
||||
) -> np.ndarray:
|
||||
raw = normalizer.inverse_transform(target_norm)
|
||||
return _decode_raw_targets(raw, pre_E)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def collect_samples(
|
||||
bundle: ModelBundle,
|
||||
val_loader: DataLoader,
|
||||
n_batches: int | None = None,
|
||||
steps: int | None = None,
|
||||
) -> SampleCollection:
|
||||
"""Run the sampler over `val_loader`, pairing generations with real targets + conditioning.
|
||||
|
||||
`steps` overrides the number of sampler steps (flow ODE steps or DDIM
|
||||
substeps); `None` keeps each sampler's own default. Unused in "ddpm"
|
||||
mode, which always runs the full schedule.
|
||||
"""
|
||||
model = bundle.model
|
||||
device = bundle.device
|
||||
steps_kw = {} if steps is None else {"steps": steps}
|
||||
|
||||
cond_list, real_list, gen_list = [], [], []
|
||||
for i, (cond_cont, cond_cat, x1, _n_sec, _sec_cont, _sec_pdg) in enumerate(val_loader):
|
||||
if n_batches is not None and i >= n_batches:
|
||||
break
|
||||
cond_cont = cond_cont.to(device)
|
||||
cond_cat = cond_cat.to(device)
|
||||
|
||||
if bundle.mode == "flow":
|
||||
gen = sample_flow(model, cond_cont, cond_cat, **steps_kw)
|
||||
elif bundle.mode == "ddpm":
|
||||
gen = sample_ddpm(model, cond_cont, cond_cat, bundle.schedule)
|
||||
else:
|
||||
gen = sample_ddim(model, cond_cont, cond_cat, bundle.schedule, **steps_kw)
|
||||
|
||||
cond_full = torch.cat([cond_cont.cpu(), cond_cat.cpu().float()], dim=-1)
|
||||
cond_list.append(cond_full.numpy())
|
||||
real_list.append(x1.numpy())
|
||||
gen_list.append(gen.cpu().numpy())
|
||||
|
||||
cond_all = np.concatenate(cond_list, axis=0)
|
||||
real_norm = np.concatenate(real_list, axis=0)
|
||||
gen_norm = np.concatenate(gen_list, axis=0)
|
||||
|
||||
cond_cont_norm, cond_cat_arr = cond_all[:, :-2], cond_all[:, -2:]
|
||||
cond_cont_raw = bundle.cond_normalizer.inverse_transform(cond_cont_norm)
|
||||
cond_cont_raw[:, 3] = inv_log_transform(cond_cont_raw[:, 3]) # log(pre_E) -> pre_E
|
||||
|
||||
idx_to_pdg, idx_to_mat = bundle.idx_to_pdg, bundle.idx_to_mat
|
||||
pdg = np.array([idx_to_pdg[i] for i in cond_cat_arr[:, 0].astype(np.int64)])
|
||||
material = np.array([idx_to_mat[i] for i in cond_cat_arr[:, 1].astype(np.int64)])
|
||||
|
||||
return SampleCollection(
|
||||
cond_cont_raw=cond_cont_raw,
|
||||
pdg=pdg,
|
||||
material=material,
|
||||
real_raw=_to_raw_targets(
|
||||
real_norm, bundle.target_normalizer, cond_cont_raw[:, 3]
|
||||
),
|
||||
gen_raw=_to_raw_targets(
|
||||
gen_norm, bundle.target_normalizer, cond_cont_raw[:, 3]
|
||||
),
|
||||
real_norm=real_norm,
|
||||
gen_norm=gen_norm,
|
||||
)
|
||||
|
||||
|
||||
def _check_predict_metadata(path: Path) -> None:
|
||||
@@ -404,39 +224,53 @@ _COND_CONT_COLS = [
|
||||
|
||||
|
||||
def load_predicted_local(
|
||||
path: str | Path, sample_frac: float = 1.0, seed: int = 0
|
||||
path: str | Path,
|
||||
sample_frac: float = 1.0,
|
||||
seed: int = 0,
|
||||
batch_size: int = 1_000_000,
|
||||
) -> SampleCollection:
|
||||
"""Build a SampleCollection from a `giant predict --coord local` parquet file.
|
||||
|
||||
Reads the `pred_*`/`true_*` columns directly — no checkpoint or model needed,
|
||||
since the predict CLI already denormalized them into the same log-scaled,
|
||||
local-frame space `collect_samples` produces internally before raw conversion.
|
||||
Requires the file to carry the `giant predict` metadata tag (see
|
||||
`_check_predict_metadata`); raises if it's missing or from --coord global,
|
||||
rather than guessing from column names.
|
||||
since the predict CLI already denormalized them into this log-scaled,
|
||||
local-frame space. Requires the file to carry the `giant predict` metadata
|
||||
tag (see `_check_predict_metadata`); raises if it's missing or from
|
||||
--coord global, rather than guessing from column names.
|
||||
|
||||
Reads via a lazy polars scan with the needed columns selected before
|
||||
`.collect()`, so column projection is pushed down into the parquet reader
|
||||
(e.g. `event_id` is never read) instead of materializing every column of
|
||||
the file as a pandas DataFrame first.
|
||||
Reads via `_iter_predicted_local_batches` (pyarrow's row-batch reader)
|
||||
with the needed columns selected, so at most `batch_size` rows are ever
|
||||
materialized at once — a plain lazy-polars `.collect()` with the sampling
|
||||
filter applied afterward looks lazy but doesn't push the row reduction
|
||||
into the scan (see its `.explain()`), so it still peaks at the full
|
||||
file's memory footprint even for a small `sample_frac`; streaming keeps
|
||||
peak memory to one batch regardless of file size or `sample_frac`.
|
||||
|
||||
`sample_frac` (0 < sample_frac <= 1) randomly keeps only that fraction of
|
||||
rows after the column projection — useful for files too large to
|
||||
comfortably hold as numpy arrays in `real_raw`/`gen_raw`. Sampling happens
|
||||
after `.collect()` since polars' row-level sampling isn't pushed down into
|
||||
the lazy scan; `seed` makes the subsample reproducible.
|
||||
comfortably hold as numpy arrays in `real_raw`/`gen_raw`. The kept/dropped
|
||||
decision is a hash of each row's position in the file (`seed`-dependent),
|
||||
computed against a running offset across batches so it's equivalent to
|
||||
hashing a single row index over the whole file rather than restarting at
|
||||
each batch boundary.
|
||||
"""
|
||||
if not (0 < sample_frac <= 1):
|
||||
raise ValueError(f"sample_frac must be in (0, 1], got {sample_frac}")
|
||||
_check_predict_metadata(Path(path))
|
||||
pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
|
||||
true_cols = [f"true_{name}" for name in LOCAL_TARGET_NAMES]
|
||||
df = (
|
||||
_scan_predicted_local(path)
|
||||
.select(pred_cols + true_cols + _COND_CONT_COLS + ["pdg", "material"])
|
||||
.collect()
|
||||
)
|
||||
if sample_frac < 1.0:
|
||||
df = df.sample(fraction=sample_frac, seed=seed)
|
||||
columns = pred_cols + true_cols + _COND_CONT_COLS + ["pdg", "material"]
|
||||
|
||||
threshold = int(sample_frac * 2**32) if sample_frac < 1.0 else None
|
||||
frames = []
|
||||
offset = 0
|
||||
for batch_df in _iter_predicted_local_batches(path, columns, batch_size):
|
||||
n = batch_df.height
|
||||
if threshold is not None:
|
||||
row_idx = pl.arange(offset, offset + n, eager=True).cast(pl.UInt32)
|
||||
batch_df = batch_df.filter((row_idx.hash(seed=seed) % 2**32) < threshold)
|
||||
offset += n
|
||||
frames.append(batch_df)
|
||||
df = pl.concat(frames) if len(frames) != 1 else frames[0]
|
||||
|
||||
gen_log_local = df.select(pred_cols).to_numpy().astype(np.float32)
|
||||
real_log_local = df.select(true_cols).to_numpy().astype(np.float32)
|
||||
@@ -527,27 +361,32 @@ def marginal_table(
|
||||
#
|
||||
# `marginal_table`/`constraint_report` above require a `SampleCollection`
|
||||
# with the full real/gen arrays already materialized in numpy. The functions
|
||||
# below instead take a parquet path (or LazyFrame) directly and stay lazy
|
||||
# end to end: each (group, dim) pair is filtered, projected to just the two
|
||||
# columns it needs, and collected on its own, so peak memory is one column
|
||||
# pair rather than the whole file — useful when `load_predicted_local` itself
|
||||
# would be too large to hold in memory at once.
|
||||
# below instead take a parquet path (or LazyFrame) directly and stay lazy,
|
||||
# reading the file a small constant number of times — one native `group_by`
|
||||
# (covering every group and every dim's mean/std/n/histogram-range at once)
|
||||
# plus one more pass per dim for histogram bin counts — rather than filtering
|
||||
# and re-collecting once per (group, dim) pair. The latter (the previous
|
||||
# implementation) meant runtime scaled with the number of *groups*: on a
|
||||
# 114M-row file with 138 distinct pdg codes, `group_by="pdg"` extrapolated to
|
||||
# roughly an hour, against ~1 minute for `group_by=None`. This version's
|
||||
# runtime is independent of group cardinality.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _histogram_kl_pl(
|
||||
p: pl.Series, q: pl.Series, bins: int = 50, eps: float = 1e-8
|
||||
def _kl_from_counts(
|
||||
real_counts: np.ndarray, gen_counts: np.ndarray, eps: float = 1e-8
|
||||
) -> float:
|
||||
"""Polars duplicate of `giant.validate._histogram_kl`, binning via `Series.hist`."""
|
||||
lo, hi = min(p.min(), q.min()), max(p.max(), q.max())
|
||||
if hi <= lo:
|
||||
return 0.0
|
||||
edges = np.linspace(lo, hi, bins + 1).tolist()
|
||||
p_hist = p.hist(bins=edges)["count"].to_numpy().astype(np.float64) + eps
|
||||
q_hist = q.hist(bins=edges)["count"].to_numpy().astype(np.float64) + eps
|
||||
p_hist /= p_hist.sum()
|
||||
q_hist /= q_hist.sum()
|
||||
return float(np.sum(p_hist * np.log(p_hist / q_hist)))
|
||||
"""KL(real || gen) from two aligned histogram bin-count arrays.
|
||||
|
||||
Same smoothing/normalization as `giant.validate._histogram_kl`, just
|
||||
taking counts directly instead of raw samples (the counts here come from
|
||||
a polars `group_by` aggregation, not `np.histogram`).
|
||||
"""
|
||||
p = real_counts.astype(np.float64) + eps
|
||||
q = gen_counts.astype(np.float64) + eps
|
||||
p /= p.sum()
|
||||
q /= q.sum()
|
||||
return float(np.sum(p * np.log(p / q)))
|
||||
|
||||
|
||||
def _raw_dim_expr(prefix: str, j: int) -> pl.Expr:
|
||||
@@ -574,33 +413,155 @@ def _scan_predicted_local(source: str | Path | pl.LazyFrame) -> pl.LazyFrame:
|
||||
return pl.scan_parquet(path)
|
||||
|
||||
|
||||
def _group_filters_pl(
|
||||
lf: pl.LazyFrame,
|
||||
group_by: str | None,
|
||||
n_energy_bins: int,
|
||||
) -> list[tuple[str, pl.Expr]]:
|
||||
def _add_group_label(
|
||||
lf: pl.LazyFrame, group_by: str | None, n_energy_bins: int
|
||||
) -> pl.LazyFrame:
|
||||
"""Add a `_group` string column matching `_group_labels`'s label format.
|
||||
|
||||
For "pdg"/"material" the label is a direct string expr over the existing
|
||||
column — no upfront pass to enumerate distinct values needed, since
|
||||
`group_by("_group")` downstream discovers them itself. "energy" still
|
||||
needs one pass over `pre_E` to fix quantile bin edges before a label can
|
||||
be assigned per row.
|
||||
"""
|
||||
if group_by is None:
|
||||
return [("all", pl.lit(True))]
|
||||
return lf.with_columns(pl.lit("all").alias("_group"))
|
||||
if group_by == "pdg":
|
||||
vals = lf.select("pdg").unique().collect()["pdg"].sort().to_list()
|
||||
return [(f"pdg={v}", pl.col("pdg") == v) for v in vals]
|
||||
return lf.with_columns(
|
||||
(pl.lit("pdg=") + pl.col("pdg").cast(pl.Int64).cast(pl.Utf8)).alias(
|
||||
"_group"
|
||||
)
|
||||
)
|
||||
if group_by == "material":
|
||||
vals = lf.select("material").unique().collect()["material"].sort().to_list()
|
||||
return [(f"material={v}", pl.col("material") == v) for v in vals]
|
||||
return lf.with_columns(
|
||||
(pl.lit("material=") + pl.col("material")).alias("_group")
|
||||
)
|
||||
if group_by == "energy":
|
||||
pre_E = lf.select("pre_E").collect().to_series().to_numpy()
|
||||
pre_E = lf.select("pre_E").collect(engine="streaming").to_series().to_numpy()
|
||||
edges = np.quantile(pre_E, np.linspace(0, 1, n_energy_bins + 1))
|
||||
edges[-1] += 1e-6
|
||||
return [
|
||||
(
|
||||
f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})",
|
||||
(pl.col("pre_E") >= edges[i]) & (pl.col("pre_E") < edges[i + 1]),
|
||||
)
|
||||
for i in range(n_energy_bins)
|
||||
labels = [
|
||||
f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})" for i in range(n_energy_bins)
|
||||
]
|
||||
expr = pl.when(pl.col("pre_E") < edges[1]).then(pl.lit(labels[0]))
|
||||
for i in range(1, n_energy_bins - 1):
|
||||
expr = expr.when(pl.col("pre_E") < edges[i + 1]).then(pl.lit(labels[i]))
|
||||
expr = expr.otherwise(pl.lit(labels[-1]))
|
||||
return lf.with_columns(expr.alias("_group"))
|
||||
raise ValueError(f"unknown group_by={group_by!r}")
|
||||
|
||||
|
||||
def _dim_narrow_lf(lf: pl.LazyFrame, j: int) -> pl.LazyFrame:
|
||||
"""Project down to just `(_group, real, gen)` for one dim, before any join/group_by.
|
||||
|
||||
Polars' projection pushdown doesn't reliably prune columns across a
|
||||
`.join()` in this version — without this explicit `.select()` up front,
|
||||
`_dim_hist_counts`'s join ends up materializing every column of the file
|
||||
(~30 columns × 114M rows ≈ 13GB) instead of just the ~2-3 this dim needs,
|
||||
even though nothing downstream references the others. Selecting first
|
||||
guarantees the narrow projection regardless of what the optimizer would
|
||||
otherwise infer.
|
||||
"""
|
||||
return lf.select(
|
||||
"_group",
|
||||
_raw_dim_expr("true", j).alias("real"),
|
||||
_raw_dim_expr("pred", j).alias("gen"),
|
||||
)
|
||||
|
||||
|
||||
def _dim_stats(narrow: pl.LazyFrame) -> pl.DataFrame:
|
||||
"""Per-group mean/std/n plus histogram-range (min/max of real+gen) for one dim.
|
||||
|
||||
A single `group_by("_group")` pass over `narrow` (see `_dim_narrow_lf`),
|
||||
mirroring the proven-cheap shape of `pdg_contribution_table_pl`, rather
|
||||
than folding all 9 dims' source columns into one query — the latter reads
|
||||
~19 columns and builds every dim's intermediate (softmax, log, etc.)
|
||||
arrays for the whole file at once, which is memory-heavy enough to OOM on
|
||||
a 114M-row file even though it's still just one pass.
|
||||
"""
|
||||
return (
|
||||
narrow.group_by("_group")
|
||||
.agg(
|
||||
pl.len().alias("n"),
|
||||
pl.col("real").mean().alias("real_mean"),
|
||||
pl.col("gen").mean().alias("gen_mean"),
|
||||
# ddof=0 to match numpy's (population-std) default used by marginal_table
|
||||
pl.col("real").std(ddof=0).alias("real_std"),
|
||||
pl.col("gen").std(ddof=0).alias("gen_std"),
|
||||
pl.min_horizontal(pl.col("real").min(), pl.col("gen").min()).alias("lo"),
|
||||
pl.max_horizontal(pl.col("real").max(), pl.col("gen").max()).alias("hi"),
|
||||
)
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
|
||||
|
||||
def _pad_degenerate_range(lo_hi: pl.DataFrame) -> pl.DataFrame:
|
||||
"""Widen a (lo, hi) pair by ±0.5 when too tight to support `bins` distinct edges.
|
||||
|
||||
Polars duplicate of `_hist_edges`'s degenerate-range handling, applied
|
||||
per group row instead of per call.
|
||||
"""
|
||||
ok = (pl.col("hi") - pl.col("lo")) > 1e-6 * pl.max_horizontal(
|
||||
pl.col("hi").abs(), pl.lit(1.0)
|
||||
)
|
||||
return lo_hi.with_columns(
|
||||
pl.when(ok).then(pl.col("lo")).otherwise(pl.col("lo") - 0.5).alias("lo"),
|
||||
pl.when(ok).then(pl.col("hi")).otherwise(pl.col("hi") + 0.5).alias("hi"),
|
||||
)
|
||||
|
||||
|
||||
def _dim_hist_counts(
|
||||
narrow: pl.LazyFrame, lo_hi: pl.DataFrame, bins: int
|
||||
) -> tuple[pl.DataFrame, pl.DataFrame]:
|
||||
"""Per-group real/gen histogram bin counts for one dim, as two long tables.
|
||||
|
||||
Joins each row (from `narrow`, see `_dim_narrow_lf`) to its group's
|
||||
(lo, hi) range (from `lo_hi`, already collected and tiny — one row per
|
||||
group), bins real/gen into `[0, bins)`, and counts via a two-key
|
||||
`group_by(["_group", "bin"])` — a proper single hash-pass histogram.
|
||||
(An earlier version aggregated with one `(bin == k).sum()` expression per
|
||||
bin, i.e. `bins` separate boolean-compare-and-reduce passes over every
|
||||
row; that's O(N × bins) work — 50 bins meant ~50x more comparisons than
|
||||
necessary — and was the dominant cost, not the join.)
|
||||
"""
|
||||
width = pl.col("hi") - pl.col("lo")
|
||||
real_bin = (
|
||||
((pl.col("real") - pl.col("lo")) / width * bins)
|
||||
.floor()
|
||||
.cast(pl.Int64)
|
||||
.clip(0, bins - 1)
|
||||
)
|
||||
gen_bin = (
|
||||
((pl.col("gen") - pl.col("lo")) / width * bins)
|
||||
.floor()
|
||||
.cast(pl.Int64)
|
||||
.clip(0, bins - 1)
|
||||
)
|
||||
|
||||
joined = narrow.join(lo_hi.lazy(), on="_group")
|
||||
real_hist = (
|
||||
joined.select("_group", real_bin.alias("bin"))
|
||||
.group_by(["_group", "bin"])
|
||||
.agg(pl.len().alias("count"))
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
gen_hist = (
|
||||
joined.select("_group", gen_bin.alias("bin"))
|
||||
.group_by(["_group", "bin"])
|
||||
.agg(pl.len().alias("count"))
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
return real_hist, gen_hist
|
||||
|
||||
|
||||
def _hist_counts_by_group(hist: pl.DataFrame, bins: int) -> dict[str, np.ndarray]:
|
||||
"""Long `(_group, bin, count)` table -> `{group: dense (bins,) count array}`."""
|
||||
out: dict[str, np.ndarray] = {}
|
||||
for group, bin_idx, count in hist.iter_rows():
|
||||
out.setdefault(group, np.zeros(bins, dtype=np.int64))[bin_idx] = count
|
||||
return out
|
||||
|
||||
|
||||
def marginal_table_pl(
|
||||
source: str | Path | pl.LazyFrame,
|
||||
group_by: str | None = None,
|
||||
@@ -612,35 +573,35 @@ def marginal_table_pl(
|
||||
`source` is a path to a `giant predict --coord local` parquet file, or an
|
||||
already-built LazyFrame with the same `pred_*`/`true_*`/`pdg`/`material`/
|
||||
`pre_E` columns (e.g. for testing). Never builds a `SampleCollection` —
|
||||
see the module-level note above on why this stays lazy.
|
||||
see the module-level note above on why this stays lazy, and on why this
|
||||
reads the file three passes per dim (stats, real histogram, gen
|
||||
histogram) rather than once per (group, dim) pair.
|
||||
"""
|
||||
lf = _scan_predicted_local(source)
|
||||
lf = _add_group_label(_scan_predicted_local(source), group_by, n_energy_bins)
|
||||
|
||||
rows = []
|
||||
for label, cond in _group_filters_pl(lf, group_by, n_energy_bins):
|
||||
glf = lf.filter(cond)
|
||||
n = glf.select(pl.len()).collect().item()
|
||||
if n < 2:
|
||||
continue
|
||||
for j, name in enumerate(RAW_TARGET_NAMES):
|
||||
pair = glf.select(
|
||||
[
|
||||
_raw_dim_expr("true", j).alias("real"),
|
||||
_raw_dim_expr("pred", j).alias("gen"),
|
||||
]
|
||||
).collect()
|
||||
real_s, gen_s = pair["real"], pair["gen"]
|
||||
for j, name in enumerate(RAW_TARGET_NAMES):
|
||||
narrow = _dim_narrow_lf(lf, j)
|
||||
stats = _dim_stats(narrow)
|
||||
stats = stats.filter(pl.col("n") >= 2)
|
||||
lo_hi = _pad_degenerate_range(stats.select("_group", "lo", "hi"))
|
||||
real_hist, gen_hist = _dim_hist_counts(narrow, lo_hi, bins)
|
||||
real_by_group = _hist_counts_by_group(real_hist, bins)
|
||||
gen_by_group = _hist_counts_by_group(gen_hist, bins)
|
||||
for row in stats.iter_rows(named=True):
|
||||
group = row["_group"]
|
||||
real_counts = real_by_group.get(group, np.zeros(bins, dtype=np.int64))
|
||||
gen_counts = gen_by_group.get(group, np.zeros(bins, dtype=np.int64))
|
||||
rows.append(
|
||||
{
|
||||
"group": label,
|
||||
"group": group,
|
||||
"dim": name,
|
||||
"n": n,
|
||||
"real_mean": real_s.mean(),
|
||||
"gen_mean": gen_s.mean(),
|
||||
# ddof=0 to match numpy's (population-std) default used by marginal_table
|
||||
"real_std": real_s.std(ddof=0),
|
||||
"gen_std": gen_s.std(ddof=0),
|
||||
"kl_real_gen": _histogram_kl_pl(real_s, gen_s, bins=bins),
|
||||
"n": row["n"],
|
||||
"real_mean": row["real_mean"],
|
||||
"gen_mean": row["gen_mean"],
|
||||
"real_std": row["real_std"],
|
||||
"gen_std": row["gen_std"],
|
||||
"kl_real_gen": _kl_from_counts(real_counts, gen_counts),
|
||||
}
|
||||
)
|
||||
return pl.DataFrame(rows).sort("kl_real_gen", descending=True)
|
||||
@@ -893,7 +854,7 @@ def direction_alignment(collection: SampleCollection) -> tuple[np.ndarray, np.nd
|
||||
def plot_direction_alignment(collection: SampleCollection, bins: int = 50):
|
||||
real_cos, gen_cos = direction_alignment(collection)
|
||||
fig, ax = plt.subplots(figsize=(5, 4))
|
||||
edges = np.linspace(-1, 1, bins + 1)
|
||||
edges = np.linspace(-1, 1, bins + 1).tolist()
|
||||
ax.hist(real_cos, bins=edges, density=True, histtype="step", label="real")
|
||||
ax.hist(gen_cos, bins=edges, density=True, histtype="step", label="generated")
|
||||
ax.set_yscale("log")
|
||||
@@ -957,8 +918,12 @@ def constraint_report_pl(
|
||||
lf = _scan_predicted_local(source)
|
||||
pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
|
||||
|
||||
post_norm = sum(pl.col(pred_cols[k]) ** 2 for k in range(3, 6)).sqrt()
|
||||
travel_norm = sum(pl.col(pred_cols[k]) ** 2 for k in range(6, 9)).sqrt()
|
||||
post_norm = pl.sum_horizontal(
|
||||
[pl.col(pred_cols[k]) ** 2 for k in range(3, 6)]
|
||||
).sqrt()
|
||||
travel_norm = pl.sum_horizontal(
|
||||
[pl.col(pred_cols[k]) ** 2 for k in range(6, 9)]
|
||||
).sqrt()
|
||||
raw_log_dims = [_raw_dim_expr("pred", j) for j in range(_N_SCALAR_DIMS)]
|
||||
|
||||
agg = (
|
||||
@@ -986,7 +951,7 @@ def constraint_report_pl(
|
||||
],
|
||||
]
|
||||
)
|
||||
.collect()
|
||||
.collect(engine="streaming")
|
||||
.row(0, named=True)
|
||||
)
|
||||
|
||||
@@ -1104,11 +1069,11 @@ def _entry_axis_and_bin_edges(
|
||||
pl.col("depth_proxy").quantile(0.999).alias("depth_hi"),
|
||||
pl.col("transverse_proxy").quantile(0.999).alias("transverse_hi"),
|
||||
)
|
||||
.collect()
|
||||
.collect(engine="streaming")
|
||||
.row(0, named=True)
|
||||
)
|
||||
|
||||
entry_df = entry.collect().sort("event_id")
|
||||
entry_df = entry.collect(engine="streaming").sort("event_id")
|
||||
|
||||
depth_lo, depth_hi = stats["depth_lo"], stats["depth_hi"]
|
||||
if not (depth_hi - depth_lo > 1e-6 * max(abs(depth_hi), 1.0)):
|
||||
@@ -1251,19 +1216,52 @@ def compute_event_observables_pl(
|
||||
real_transverse_bin = np.digitize(real_transverse, transverse_edges[1:-1])
|
||||
gen_transverse_bin = np.digitize(gen_transverse, transverse_edges[1:-1])
|
||||
|
||||
np.add.at(n_steps, idx, 1)
|
||||
np.add.at(real_total_edep, idx, real_edep)
|
||||
np.add.at(gen_total_edep, idx, gen_edep)
|
||||
np.add.at(real_total_length, idx, real_step_length)
|
||||
np.add.at(gen_total_length, idx, gen_step_length)
|
||||
np.add.at(real_sum_edep_depth, idx, real_edep * real_depth)
|
||||
np.add.at(gen_sum_edep_depth, idx, gen_edep * gen_depth)
|
||||
np.add.at(real_sum_edep_transverse2, idx, real_edep * real_transverse**2)
|
||||
np.add.at(gen_sum_edep_transverse2, idx, gen_edep * gen_transverse**2)
|
||||
np.add.at(real_depth_bin_edep, (idx, real_depth_bin), real_edep)
|
||||
np.add.at(gen_depth_bin_edep, (idx, gen_depth_bin), gen_edep)
|
||||
np.add.at(real_transverse_bin_edep, (idx, real_transverse_bin), real_edep)
|
||||
np.add.at(gen_transverse_bin_edep, (idx, gen_transverse_bin), gen_edep)
|
||||
# np.add.at is an unbuffered ufunc method — a well-known numpy slow path
|
||||
# for scatter-add (44% of this function's runtime on a 114M-row profile).
|
||||
# np.bincount does the same accumulation with a single optimized pass;
|
||||
# the 2D (per-event, per-bin) accumulators flatten (idx, bin) into one
|
||||
# bincount index and reshape back, since bincount only scatters into 1D.
|
||||
n_steps += np.bincount(idx, minlength=n_events)
|
||||
real_total_edep += np.bincount(idx, weights=real_edep, minlength=n_events)
|
||||
gen_total_edep += np.bincount(idx, weights=gen_edep, minlength=n_events)
|
||||
real_total_length += np.bincount(
|
||||
idx, weights=real_step_length, minlength=n_events
|
||||
)
|
||||
gen_total_length += np.bincount(
|
||||
idx, weights=gen_step_length, minlength=n_events
|
||||
)
|
||||
real_sum_edep_depth += np.bincount(
|
||||
idx, weights=real_edep * real_depth, minlength=n_events
|
||||
)
|
||||
gen_sum_edep_depth += np.bincount(
|
||||
idx, weights=gen_edep * gen_depth, minlength=n_events
|
||||
)
|
||||
real_sum_edep_transverse2 += np.bincount(
|
||||
idx, weights=real_edep * real_transverse**2, minlength=n_events
|
||||
)
|
||||
gen_sum_edep_transverse2 += np.bincount(
|
||||
idx, weights=gen_edep * gen_transverse**2, minlength=n_events
|
||||
)
|
||||
real_depth_bin_edep += np.bincount(
|
||||
idx * depth_bins + real_depth_bin,
|
||||
weights=real_edep,
|
||||
minlength=n_events * depth_bins,
|
||||
).reshape(n_events, depth_bins)
|
||||
gen_depth_bin_edep += np.bincount(
|
||||
idx * depth_bins + gen_depth_bin,
|
||||
weights=gen_edep,
|
||||
minlength=n_events * depth_bins,
|
||||
).reshape(n_events, depth_bins)
|
||||
real_transverse_bin_edep += np.bincount(
|
||||
idx * transverse_bins + real_transverse_bin,
|
||||
weights=real_edep,
|
||||
minlength=n_events * transverse_bins,
|
||||
).reshape(n_events, transverse_bins)
|
||||
gen_transverse_bin_edep += np.bincount(
|
||||
idx * transverse_bins + gen_transverse_bin,
|
||||
weights=gen_edep,
|
||||
minlength=n_events * transverse_bins,
|
||||
).reshape(n_events, transverse_bins)
|
||||
|
||||
safe_real_total = np.where(real_total_edep > 0, real_total_edep, 1.0)
|
||||
safe_gen_total = np.where(gen_total_edep > 0, gen_total_edep, 1.0)
|
||||
@@ -1312,7 +1310,7 @@ def compute_event_observables_pl(
|
||||
.alias("gen_median_length"),
|
||||
)
|
||||
.sort("event_id")
|
||||
.collect()
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
real_median_edep = medians["real_median_edep"].to_numpy()
|
||||
gen_median_edep = medians["gen_median_edep"].to_numpy()
|
||||
@@ -1366,7 +1364,7 @@ def plot_total_energy(observables: EventObservables, bins: int = 50):
|
||||
gen = table["gen_total_edep"].to_numpy()
|
||||
|
||||
fig, ax = plt.subplots(figsize=(6, 4))
|
||||
edges = _hist_edges(real, gen, bins=bins)
|
||||
edges = _hist_edges(real, gen, bins=bins).tolist()
|
||||
ax.hist(
|
||||
real,
|
||||
bins=edges,
|
||||
@@ -1395,7 +1393,7 @@ def plot_total_length(observables: EventObservables, bins: int = 50):
|
||||
gen = table["gen_total_length"].to_numpy()
|
||||
|
||||
fig, ax = plt.subplots(figsize=(6, 4))
|
||||
edges = _hist_edges(real, gen, bins=bins)
|
||||
edges = _hist_edges(real, gen, bins=bins).tolist()
|
||||
ax.hist(
|
||||
real,
|
||||
bins=edges,
|
||||
@@ -1577,7 +1575,7 @@ def plot_shower_max_depth(observables: EventObservables, bins: int = 30):
|
||||
gen = table["gen_max_depth"].to_numpy()
|
||||
|
||||
fig, ax = plt.subplots(figsize=(6, 4))
|
||||
edges = _hist_edges(real, gen, bins=bins)
|
||||
edges = _hist_edges(real, gen, bins=bins).tolist()
|
||||
ax.hist(real, bins=edges, density=True, histtype="step", label="real")
|
||||
ax.hist(gen, bins=edges, density=True, histtype="step", label="generated")
|
||||
ax.set_xlabel("depth of shower maximum [mm]")
|
||||
@@ -1627,7 +1625,7 @@ def pdg_contribution_table_pl(source: str | Path | pl.LazyFrame) -> pl.DataFrame
|
||||
_delog("true_log_step_length").sum().alias("real_total_length"),
|
||||
_delog("pred_log_step_length").sum().alias("gen_total_length"),
|
||||
)
|
||||
.collect()
|
||||
.collect(engine="streaming")
|
||||
.sort("pdg")
|
||||
)
|
||||
|
||||
|
||||
+19
-7
@@ -1,5 +1,5 @@
|
||||
from collections import Counter
|
||||
from datetime import datetime, timezone
|
||||
from datetime import date, datetime, timezone
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
@@ -45,9 +45,7 @@ app = typer.Typer(no_args_is_help=True)
|
||||
_CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions")
|
||||
|
||||
|
||||
def _resolve_prediction_output(
|
||||
data: Path, out: Path | None
|
||||
) -> tuple[Path, Path, str]:
|
||||
def _resolve_prediction_output(data: Path, out: Path | None) -> tuple[Path, Path, str]:
|
||||
"""Return (out_path, resolved_dataset_path, pred_uuid).
|
||||
|
||||
When *out* is None the output path is derived from *data*:
|
||||
@@ -69,6 +67,7 @@ def _write_prediction_ref(
|
||||
pred_uuid: str,
|
||||
out: Path,
|
||||
dataset_path: Path,
|
||||
comment: str | None = None,
|
||||
) -> Path:
|
||||
"""Write a YAML sidecar in the checkpoint directory and return its path."""
|
||||
ref = {
|
||||
@@ -78,6 +77,8 @@ def _write_prediction_ref(
|
||||
"checkpoint": str(checkpoint.resolve()),
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
if comment is not None:
|
||||
ref["comment"] = comment
|
||||
ref_path = checkpoint.parent / f"{pred_uuid}.yaml"
|
||||
ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False))
|
||||
return ref_path
|
||||
@@ -245,7 +246,8 @@ def train(
|
||||
)
|
||||
|
||||
out_dir = out or Path(
|
||||
f"checkpoints/{t['mode']}"
|
||||
f"checkpoints/{date.today().strftime('%Y%m%d')}"
|
||||
f"_{t['mode']}"
|
||||
f"_h{m['hidden_dim']}"
|
||||
f"_b{m['n_blocks']}"
|
||||
f"_e{m['emb_dim']}"
|
||||
@@ -316,6 +318,14 @@ def predict(
|
||||
Optional[str],
|
||||
typer.Option("--device", "-d", help="cpu | cuda | mps (default: auto)"),
|
||||
] = None,
|
||||
comment: Annotated[
|
||||
Optional[str],
|
||||
typer.Option(
|
||||
"--comment",
|
||||
"-m",
|
||||
help="Free-text note recorded in the prediction's YAML sidecar",
|
||||
),
|
||||
] = None,
|
||||
) -> None:
|
||||
"""Run trained model on a parquet file and save predictions."""
|
||||
batch_size_auto = False
|
||||
@@ -398,7 +408,7 @@ def predict(
|
||||
nonlocal writer, total
|
||||
|
||||
if coord == Coord.local:
|
||||
cond_cont, cond_cat, target_raw, _, _ = build_features(
|
||||
cond_cont, cond_cat, target_raw, _, _, _, _, _ = build_features(
|
||||
piece, pdg_map, mat_map
|
||||
)
|
||||
cond_cont = cond_norm.transform(cond_cont)
|
||||
@@ -534,7 +544,9 @@ def predict(
|
||||
if writer is not None:
|
||||
writer.close()
|
||||
|
||||
ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path)
|
||||
ref_path = _write_prediction_ref(
|
||||
checkpoint, pred_uuid, out, dataset_path, comment
|
||||
)
|
||||
typer.echo(f"reference: {ref_path}")
|
||||
|
||||
if skipped:
|
||||
|
||||
+76
-23
@@ -1,3 +1,5 @@
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
|
||||
_EPS = 1e-8
|
||||
@@ -37,10 +39,26 @@ def energy_simplex_encode(
|
||||
`energy_simplex_decode` is its inverse (up to the floor softening).
|
||||
"""
|
||||
pre_E = np.maximum(np.asarray(pre_E, dtype=np.float32), _EPS)
|
||||
post_E = np.clip(np.asarray(post_E, dtype=np.float32), 0.0, pre_E)
|
||||
raw_post_E = np.asarray(post_E, dtype=np.float32)
|
||||
post_E = np.clip(raw_post_E, 0.0, pre_E)
|
||||
delta_e = pre_E - post_E
|
||||
lost = np.asarray(edep, dtype=np.float32) + np.asarray(e_sec, dtype=np.float32)
|
||||
edep = np.asarray(edep, dtype=np.float32)
|
||||
e_sec = np.asarray(e_sec, dtype=np.float32)
|
||||
lost = edep + e_sec
|
||||
has_loss = lost > _EPS
|
||||
# Clipping post_E down to pre_E forces delta_e (and thus the rescaled
|
||||
# edep/e_sec below) to 0 even on rows where edep/e_sec were genuinely
|
||||
# recorded as nonzero — warn so this doesn't silently discard real data.
|
||||
discarded = (raw_post_E > pre_E) & has_loss
|
||||
if np.any(discarded):
|
||||
n = int(np.sum(discarded))
|
||||
warnings.warn(
|
||||
f"energy_simplex_encode: {n}/{len(discarded)} step(s) had "
|
||||
"post_E > pre_E (clipped) while recording nonzero edep/e_sec; "
|
||||
"that recorded energy deposit is discarded to keep delta_e "
|
||||
"consistent with the clip.",
|
||||
stacklevel=2,
|
||||
)
|
||||
scale = np.where(has_loss, delta_e / np.maximum(lost, _EPS), 0.0)
|
||||
# Where nothing was recorded as deposited/secondary but energy was lost,
|
||||
# attribute all of delta_e to local deposit.
|
||||
@@ -80,6 +98,54 @@ def energy_simplex_decode(
|
||||
)
|
||||
|
||||
|
||||
def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
|
||||
"""Unit rotation axis `pre_dir × ẑ`, closed-form since ẑ = [0, 0, 1] is constant.
|
||||
|
||||
`cross(a, [0,0,1]) = [a_y, -a_x, 0]` — substituting the constant operand
|
||||
avoids a generic `np.cross` call (shape/broadcast handling for an
|
||||
arbitrary second operand) on every row; profiling on a 114M-row file
|
||||
showed `np.cross` as the single hottest call inside this rotation.
|
||||
"""
|
||||
axis = np.stack(
|
||||
[pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1
|
||||
)
|
||||
axis_norm = np.linalg.norm(axis, axis=1, keepdims=True) # (N,1)
|
||||
# Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis
|
||||
# are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant).
|
||||
safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
|
||||
return np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)
|
||||
|
||||
|
||||
def _cross_with_z_axis(axis: np.ndarray, v: np.ndarray) -> np.ndarray:
|
||||
"""`axis × v`, closed-form since `axis` from `_rodrigues_axis` always has z = 0.
|
||||
|
||||
`cross([ax,ay,0], [bx,by,bz]) = [ay*bz, -ax*bz, ax*by - ay*bx]`.
|
||||
"""
|
||||
ax, ay = axis[:, 0:1], axis[:, 1:2]
|
||||
bx, by, bz = v[:, 0:1], v[:, 1:2], v[:, 2:3]
|
||||
return np.concatenate([ay * bz, -ax * bz, ax * by - ay * bx], axis=1)
|
||||
|
||||
|
||||
def _validate_unit_pre_dir(pre_dir: np.ndarray) -> np.ndarray:
|
||||
"""Normalize pre_dir and raise if any row is too degenerate to define a frame.
|
||||
|
||||
`local_frame_rotation`/`inv_local_frame_rotation` treat pre_dir[:, 2] as
|
||||
cos(angle to ẑ), which is only correct for a unit vector. Small float32
|
||||
drift is corrected silently; a near-zero-norm row has no well-defined
|
||||
direction, so it's raised loudly instead of producing a meaningless
|
||||
rotation (previously it fell through to an arbitrary axis with no error).
|
||||
"""
|
||||
pre_dir = np.asarray(pre_dir, dtype=np.float32)
|
||||
norm = np.linalg.norm(pre_dir, axis=1, keepdims=True)
|
||||
if np.any(norm < 1e-6):
|
||||
raise ValueError(
|
||||
f"pre_dir has {int(np.sum(norm < 1e-6))} row(s) with near-zero norm "
|
||||
"(< 1e-6); local/inv_local_frame_rotation require a well-defined "
|
||||
"incoming direction for every row."
|
||||
)
|
||||
return pre_dir / norm
|
||||
|
||||
|
||||
def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarray:
|
||||
"""Rotate post_dir into the local frame where pre_dir maps to ẑ (Rodrigues).
|
||||
|
||||
@@ -87,20 +153,12 @@ def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarra
|
||||
expressed relative to a coordinate system in which the incoming particle
|
||||
travels along +z.
|
||||
"""
|
||||
z = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
cos_t = np.clip((pre_dir * z).sum(axis=1, keepdims=True), -1.0, 1.0) # (N,1)
|
||||
pre_dir = _validate_unit_pre_dir(pre_dir)
|
||||
cos_t = np.clip(pre_dir[:, 2:3], -1.0, 1.0) # (N,1); dot with ẑ = z-component
|
||||
sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t**2)) # (N,1)
|
||||
|
||||
axis = np.cross(pre_dir, z) # (N,3); zero when pre_dir ∥ ẑ
|
||||
axis_norm = np.linalg.norm(axis, axis=1, keepdims=True) # (N,1)
|
||||
|
||||
# Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis
|
||||
# are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant).
|
||||
safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
|
||||
axis = np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)
|
||||
|
||||
kxv = np.cross(axis, post_dir) # (N,3)
|
||||
axis = _rodrigues_axis(pre_dir) # (N,3); zero-z, zero-norm when pre_dir ∥ ẑ
|
||||
kxv = _cross_with_z_axis(axis, post_dir) # (N,3)
|
||||
kdv = (axis * post_dir).sum(axis=1, keepdims=True) # (N,1)
|
||||
|
||||
return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
|
||||
@@ -208,17 +266,12 @@ def inv_local_frame_rotation(
|
||||
|
||||
Applies R^T (same axis, negative angle) to post_dir_local.
|
||||
"""
|
||||
z = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
cos_t = np.clip((pre_dir * z).sum(axis=1, keepdims=True), -1.0, 1.0)
|
||||
pre_dir = _validate_unit_pre_dir(pre_dir)
|
||||
cos_t = np.clip(pre_dir[:, 2:3], -1.0, 1.0) # dot with ẑ = z-component
|
||||
sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t**2))
|
||||
|
||||
axis = np.cross(pre_dir, z)
|
||||
axis_norm = np.linalg.norm(axis, axis=1, keepdims=True)
|
||||
safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
|
||||
axis = np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)
|
||||
|
||||
kxv = np.cross(axis, post_dir_local)
|
||||
axis = _rodrigues_axis(pre_dir)
|
||||
kxv = _cross_with_z_axis(axis, post_dir_local)
|
||||
kdv = (axis * post_dir_local).sum(axis=1, keepdims=True)
|
||||
|
||||
# Negative angle: sin_t → -sin_t
|
||||
|
||||
+21
-2
@@ -198,9 +198,28 @@ def train(
|
||||
start_epoch = ckpt.get("epoch", 0) + 1
|
||||
best_val_loss = ckpt.get("best_val_loss", float("inf"))
|
||||
|
||||
# optimizer/lr_sched.load_state_dict() above restore the checkpoint's
|
||||
# own base LR, which would otherwise silently override an explicit
|
||||
# `lr` argument. Make `lr` authoritative again, applied at whatever
|
||||
# point the cosine/warmup schedule has already reached.
|
||||
lr_sched.base_lrs = [lr for _ in lr_sched.base_lrs]
|
||||
resumed_lr = lr * _lr_lambda(lr_sched.last_epoch)
|
||||
for group in optimizer.param_groups:
|
||||
group["lr"] = resumed_lr
|
||||
|
||||
if start_epoch > epochs:
|
||||
print(
|
||||
f"checkpoint already completed epoch {start_epoch - 1} "
|
||||
f"(>= --epochs {epochs}) — nothing to train"
|
||||
)
|
||||
return
|
||||
|
||||
metrics_path = out_dir / "metrics.csv"
|
||||
write_header = not (resume_path is not None and metrics_path.exists())
|
||||
metrics_file = open(metrics_path, "a", newline="")
|
||||
resuming_existing_metrics = resume_path is not None and metrics_path.exists()
|
||||
write_header = not resuming_existing_metrics
|
||||
metrics_file = open(
|
||||
metrics_path, "a" if resuming_existing_metrics else "w", newline=""
|
||||
)
|
||||
metrics_writer = csv.DictWriter(metrics_file, fieldnames=_METRICS_FIELDS)
|
||||
if write_header:
|
||||
metrics_writer.writeheader()
|
||||
|
||||
+2
-5
@@ -24,6 +24,7 @@ dev = [
|
||||
"pytest>=8,<10",
|
||||
"ruff>=0.15,<1",
|
||||
"ty>=0.0.50,<0.1",
|
||||
"giant[convert,analysis]",
|
||||
]
|
||||
convert = [
|
||||
"uproot>=5.3,<6",
|
||||
@@ -38,11 +39,7 @@ analysis = [
|
||||
|
||||
[project.scripts]
|
||||
giant = "giant.cli:app"
|
||||
steps-to-parquet = "scripts.steps_to_parquet:main"
|
||||
steps-to-parquet-parallel = "scripts.steps_to_parquet_parallel:main"
|
||||
migrate-geant-steps = "scripts.migrate_geant_steps:main"
|
||||
bump-dataset-version = "scripts.bump_dataset_version:main"
|
||||
create-root-files = "scripts.create_root_files:main"
|
||||
dwarf = "scripts.dwarf:app"
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
|
||||
+526
-223
@@ -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.")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
+14
-25
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import torch
|
||||
import torch.version
|
||||
|
||||
print(f"PyTorch version: {torch.__version__}")
|
||||
print(f"CUDA available: {torch.cuda.is_available()}")
|
||||
|
||||
@@ -84,8 +84,6 @@ def _make_collection(n=200, seed=0, gen_offset=0.0) -> SampleCollection:
|
||||
material=rng.choice(["W", "Pb"], size=n),
|
||||
real_raw=real,
|
||||
gen_raw=gen,
|
||||
real_norm=real,
|
||||
gen_norm=gen,
|
||||
)
|
||||
|
||||
|
||||
@@ -259,8 +257,6 @@ def test_load_predicted_local_round_trips_values(tmp_path):
|
||||
np.testing.assert_allclose(
|
||||
collection.gen_raw, expected_raw(pred_log_local), atol=1e-4
|
||||
)
|
||||
assert collection.real_norm is None
|
||||
assert collection.gen_norm is None
|
||||
|
||||
|
||||
def test_load_predicted_local_usable_by_downstream_plots(tmp_path):
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
import os
|
||||
import pytest
|
||||
from scripts import bump_dataset_version
|
||||
|
||||
plan_bump_gen = bump_dataset_version.plan_bump_gen
|
||||
@@ -13,7 +12,9 @@ check_holdout_overlap = bump_dataset_version.check_holdout_overlap
|
||||
|
||||
|
||||
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")
|
||||
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",
|
||||
@@ -55,14 +56,18 @@ def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
|
||||
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")
|
||||
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")
|
||||
dirs, _ = plan_bump_schema(
|
||||
tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01"
|
||||
)
|
||||
assert dirs == [tmp_path / "processed" / "steps" / "gen2" / "schema1"]
|
||||
|
||||
|
||||
@@ -74,6 +79,50 @@ def test_bump_schema_rejects_nonexistent_gen(tmp_path):
|
||||
pass
|
||||
|
||||
|
||||
def test_bump_gen_to_specific_tag(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
dirs, log_line = plan_bump_gen(
|
||||
tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5"
|
||||
)
|
||||
assert dirs[0] == tmp_path / "raw" / "steps" / "gen5"
|
||||
assert "`gen5`" in log_line
|
||||
|
||||
|
||||
def test_bump_gen_rejects_invalid_to_tag(tmp_path):
|
||||
try:
|
||||
plan_bump_gen(tmp_path, "steps", "bad tag", None, "2026-01-01", target="v5")
|
||||
assert False, "expected SystemExit"
|
||||
except SystemExit:
|
||||
pass
|
||||
|
||||
|
||||
def test_bump_schema_to_specific_tag(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
(tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True)
|
||||
dirs, log_line = plan_bump_schema(
|
||||
tmp_path,
|
||||
"steps",
|
||||
"gen1",
|
||||
"jump to schema5",
|
||||
None,
|
||||
"2026-01-01",
|
||||
target="schema5",
|
||||
)
|
||||
assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema5"]
|
||||
assert "`schema5`" in log_line
|
||||
|
||||
|
||||
def test_bump_schema_rejects_invalid_to_tag(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
try:
|
||||
plan_bump_schema(
|
||||
tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3"
|
||||
)
|
||||
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)
|
||||
@@ -97,6 +146,7 @@ def test_apply_bump_appends_without_clobbering_existing_log(tmp_path):
|
||||
# update-manifest
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_parquet(path):
|
||||
"""Create a zero-byte stand-in for a parquet file."""
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
@@ -104,7 +154,15 @@ def _make_parquet(path):
|
||||
|
||||
|
||||
def test_update_manifest_bumps_to_specified_schema(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema2"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
@@ -126,7 +184,15 @@ def test_update_manifest_auto_detects_highest_schema(tmp_path):
|
||||
for schema in ("schema1", "schema2", "schema3"):
|
||||
d = tmp_path / "processed" / "steps" / "gen1" / schema / "pbwo4"
|
||||
d.mkdir(parents=True)
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema3"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
parquet.touch()
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
@@ -155,7 +221,15 @@ def test_update_manifest_reports_missing_targets(tmp_path):
|
||||
|
||||
|
||||
def test_update_manifest_skips_already_at_target(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema2"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
@@ -170,7 +244,15 @@ def test_update_manifest_skips_already_at_target(tmp_path):
|
||||
|
||||
|
||||
def test_update_manifest_preserves_comments_and_blanks(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema2"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
@@ -185,8 +267,66 @@ def test_update_manifest_preserves_comments_and_blanks(tmp_path):
|
||||
assert lines[2][1] is not None # the data line was updated
|
||||
|
||||
|
||||
def test_update_manifest_bumps_gen(tmp_path):
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen2"
|
||||
/ "schema1"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
manifest_dir.mkdir(parents=True)
|
||||
manifest = manifest_dir / "full.manifest"
|
||||
manifest.write_text("../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n")
|
||||
|
||||
lines, missing = plan_update_manifest(manifest, None, target_gen="gen2")
|
||||
assert missing == []
|
||||
changed = [(old, new) for old, new in lines if new is not None]
|
||||
assert len(changed) == 1
|
||||
assert "gen2" in changed[0][1]
|
||||
assert "gen1" not in changed[0][1]
|
||||
|
||||
|
||||
def test_update_manifest_bumps_gen_and_schema(tmp_path):
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen2"
|
||||
/ "schema3"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
manifest_dir.mkdir(parents=True)
|
||||
manifest = manifest_dir / "full.manifest"
|
||||
manifest.write_text("../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n")
|
||||
|
||||
lines, missing = plan_update_manifest(manifest, "schema3", target_gen="gen2")
|
||||
assert missing == []
|
||||
changed = [(old, new) for old, new in lines if new is not None]
|
||||
assert len(changed) == 1
|
||||
assert "gen2" in changed[0][1]
|
||||
assert "schema3" in changed[0][1]
|
||||
|
||||
|
||||
def test_apply_update_manifest_writes_file(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema2"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
@@ -206,9 +346,26 @@ def test_apply_update_manifest_writes_file(tmp_path):
|
||||
# create-manifest
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_create_manifest_writes_relative_paths(tmp_path):
|
||||
pq1 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
pq2 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-001.parquet"
|
||||
pq1 = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema2"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
pq2 = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema2"
|
||||
/ "pbwo4"
|
||||
/ "shard-001.parquet"
|
||||
)
|
||||
_make_parquet(pq1)
|
||||
_make_parquet(pq2)
|
||||
|
||||
@@ -217,8 +374,8 @@ def test_create_manifest_writes_relative_paths(tmp_path):
|
||||
|
||||
assert missing == []
|
||||
assert len(lines) == 2
|
||||
assert all("schema2" in l for l in lines)
|
||||
assert all(not l.startswith("/") for l in lines)
|
||||
assert all("schema2" in line for line in lines)
|
||||
assert all(not line.startswith("/") for line in lines)
|
||||
assert resolved == [pq1.resolve(), pq2.resolve()]
|
||||
|
||||
apply_create_manifest(output, lines)
|
||||
@@ -248,6 +405,7 @@ def test_create_manifest_creates_parent_dirs(tmp_path):
|
||||
# check_holdout_overlap
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_no_overlap_check_when_no_holdout_involved(tmp_path):
|
||||
pool_dir = tmp_path / "pools" / "pbwo4"
|
||||
pool_dir.mkdir(parents=True)
|
||||
|
||||
@@ -2,7 +2,11 @@ import uuid
|
||||
|
||||
import yaml
|
||||
|
||||
from giant.cli import _CEPH_PREDICTIONS, _resolve_prediction_output, _write_prediction_ref
|
||||
from giant.cli import (
|
||||
_CEPH_PREDICTIONS,
|
||||
_resolve_prediction_output,
|
||||
_write_prediction_ref,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -99,6 +103,25 @@ def test_ref_yaml_contains_expected_fields(tmp_path):
|
||||
assert data["dataset"] == str(dataset)
|
||||
assert data["checkpoint"] == str(checkpoint.resolve())
|
||||
assert "timestamp" in data
|
||||
assert "comment" not in data
|
||||
|
||||
|
||||
def test_ref_yaml_includes_comment_when_provided(tmp_path):
|
||||
ckpt_dir = tmp_path / "checkpoints"
|
||||
ckpt_dir.mkdir()
|
||||
checkpoint = ckpt_dir / "best.pt"
|
||||
checkpoint.touch()
|
||||
|
||||
out = tmp_path / "pred.parquet"
|
||||
dataset = tmp_path / "full.manifest"
|
||||
pred_uuid = str(uuid.uuid4())
|
||||
|
||||
ref_path = _write_prediction_ref(
|
||||
checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3"
|
||||
)
|
||||
data = yaml.safe_load(ref_path.read_text())
|
||||
|
||||
assert data["comment"] == "baseline sweep run 3"
|
||||
|
||||
|
||||
def test_ref_timestamp_is_iso_format(tmp_path):
|
||||
@@ -110,7 +133,9 @@ def test_ref_timestamp_is_iso_format(tmp_path):
|
||||
checkpoint.touch()
|
||||
|
||||
pred_uuid = str(uuid.uuid4())
|
||||
ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
|
||||
ref_path = _write_prediction_ref(
|
||||
checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
|
||||
)
|
||||
data = yaml.safe_load(ref_path.read_text())
|
||||
|
||||
# Must parse without error and be timezone-aware (UTC).
|
||||
@@ -125,7 +150,9 @@ def test_ref_checkpoint_path_is_absolute(tmp_path):
|
||||
checkpoint.touch()
|
||||
|
||||
pred_uuid = str(uuid.uuid4())
|
||||
ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
|
||||
ref_path = _write_prediction_ref(
|
||||
checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
|
||||
)
|
||||
data = yaml.safe_load(ref_path.read_text())
|
||||
|
||||
assert data["checkpoint"].startswith("/")
|
||||
|
||||
@@ -50,7 +50,10 @@ sys.exit({exit_code})
|
||||
|
||||
|
||||
def test_parse_detector_spec_with_config():
|
||||
assert parse_detector_spec("sampling_pb_scint:pb_scint") == ("sampling_pb_scint", "pb_scint")
|
||||
assert parse_detector_spec("sampling_pb_scint:pb_scint") == (
|
||||
"sampling_pb_scint",
|
||||
"pb_scint",
|
||||
)
|
||||
|
||||
|
||||
def test_parse_detector_spec_without_config():
|
||||
@@ -79,13 +82,17 @@ def test_next_shard_index_continues_past_existing(tmp_path):
|
||||
|
||||
def test_plan_jobs_rejects_missing_gen(tmp_path):
|
||||
with pytest.raises(PlanError):
|
||||
plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1")
|
||||
plan_jobs(
|
||||
["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1"
|
||||
)
|
||||
|
||||
|
||||
def test_plan_jobs_rejects_malformed_gen(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
with pytest.raises(PlanError):
|
||||
plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen")
|
||||
plan_jobs(
|
||||
["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen"
|
||||
)
|
||||
|
||||
|
||||
def test_plan_jobs_continues_from_existing_shards(tmp_path):
|
||||
@@ -94,7 +101,9 @@ def test_plan_jobs_continues_from_existing_shards(tmp_path):
|
||||
(gen_dir / "pbwo4" / "shard-000.root").touch()
|
||||
(gen_dir / "pbwo4" / "shard-001.root").touch()
|
||||
|
||||
jobs = plan_jobs(["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1")
|
||||
jobs = plan_jobs(
|
||||
["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1"
|
||||
)
|
||||
|
||||
assert [j.shard_index for j in jobs] == [2, 3, 4]
|
||||
assert all(j.detector == "pbwo4" and j.config is None for j in jobs)
|
||||
@@ -133,6 +142,7 @@ def test_run_job_moves_output_to_correct_shard_path(tmp_path):
|
||||
|
||||
assert result.ok
|
||||
assert result.dest == gen_dir / "pbwo4" / "shard-007.root"
|
||||
assert result.dest is not None
|
||||
assert result.dest.is_file()
|
||||
assert not any(tmp_root.iterdir()) # workdir cleaned up
|
||||
|
||||
@@ -147,6 +157,7 @@ def test_run_job_passes_config_arg_and_isolates_cwd(tmp_path):
|
||||
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
||||
|
||||
assert result.ok
|
||||
assert result.dest is not None
|
||||
payload = json.loads(result.dest.read_text())
|
||||
assert payload["argv"] == ["pb_scint", "10000"]
|
||||
# ran in its own scratch workdir under .sim-tmp, not directly in dataset_root
|
||||
@@ -163,6 +174,7 @@ def test_run_job_omits_config_arg_when_none(tmp_path):
|
||||
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
||||
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
||||
|
||||
assert result.dest is not None
|
||||
payload = json.loads(result.dest.read_text())
|
||||
assert payload["argv"] == ["10000"]
|
||||
|
||||
@@ -229,12 +241,15 @@ def test_run_all_caps_concurrency(tmp_path):
|
||||
tmp_root.mkdir()
|
||||
|
||||
jobs = [SimJob(detector="pbwo4", config=None, shard_index=i) for i in range(6)]
|
||||
results = run_all(jobs, fake, 10000, tmp_path, "steps", "gen1", max_workers=2, tmp_root=tmp_root)
|
||||
results = run_all(
|
||||
jobs, fake, 10000, tmp_path, "steps", "gen1", max_workers=2, tmp_root=tmp_root
|
||||
)
|
||||
|
||||
assert all(r.ok for r in results)
|
||||
assert {r.dest.name for r in results} == {f"shard-{i:03d}.root" for i in range(6)}
|
||||
assert all(r.ok and r.dest is not None for r in results)
|
||||
dests = [r.dest for r in results if r.dest is not None]
|
||||
assert {d.name for d in dests} == {f"shard-{i:03d}.root" for i in range(6)}
|
||||
|
||||
intervals = [json.loads(r.dest.read_text()) for r in results]
|
||||
intervals = [json.loads(d.read_text()) for d in dests]
|
||||
events = sorted(
|
||||
[(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals]
|
||||
)
|
||||
@@ -244,24 +259,3 @@ def test_run_all_caps_concurrency(tmp_path):
|
||||
concurrent += delta
|
||||
peak = max(peak, concurrent)
|
||||
assert peak <= 2
|
||||
|
||||
|
||||
def test_build_parser_defaults():
|
||||
args = create_root_files.build_parser().parse_args(
|
||||
[
|
||||
"--executable",
|
||||
"fake",
|
||||
"--detector",
|
||||
"pbwo4",
|
||||
"--num-files",
|
||||
"2",
|
||||
"--events-per-file",
|
||||
"10000",
|
||||
"--gen",
|
||||
"gen1",
|
||||
]
|
||||
)
|
||||
assert args.jobs == 4
|
||||
assert args.kind == "steps"
|
||||
assert args.dataset_root == "/ceph/lbogner/geant_steps"
|
||||
assert args.execute is False
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from scripts.dwarf import app
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
||||
def test_convert_rejects_jobs_below_one(tmp_path):
|
||||
root_file = tmp_path / "shard.root"
|
||||
root_file.touch()
|
||||
result = runner.invoke(app, ["convert", str(root_file), "--jobs", "0"])
|
||||
assert result.exit_code != 0
|
||||
assert "--jobs must be >= 1" in result.output
|
||||
|
||||
|
||||
def test_convert_rejects_output_with_multiple_files(tmp_path):
|
||||
a = tmp_path / "a.root"
|
||||
b = tmp_path / "b.root"
|
||||
a.touch()
|
||||
b.touch()
|
||||
result = runner.invoke(app, ["convert", str(a), str(b), "--output", "out.parquet"])
|
||||
assert result.exit_code != 0
|
||||
assert "--output can only be used with a single input file" in result.output
|
||||
|
||||
|
||||
def test_convert_rejects_output_with_parallel_jobs(tmp_path):
|
||||
root_file = tmp_path / "shard.root"
|
||||
root_file.touch()
|
||||
result = runner.invoke(
|
||||
app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"]
|
||||
)
|
||||
assert result.exit_code != 0
|
||||
assert "--output cannot be combined with --jobs > 1" in result.output
|
||||
|
||||
|
||||
def test_convert_default_jobs_is_one():
|
||||
result = runner.invoke(app, ["convert", "--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "default: 1" in result.output
|
||||
|
||||
|
||||
def test_bump_gen_requires_reason():
|
||||
result = runner.invoke(app, ["bump-gen"])
|
||||
assert result.exit_code != 0
|
||||
assert "reason" in result.output.lower()
|
||||
|
||||
|
||||
def test_create_manifest_requires_exactly_one_of_output_or_pool(tmp_path):
|
||||
f = tmp_path / "a.parquet"
|
||||
f.touch()
|
||||
result = runner.invoke(app, ["create-manifest", str(f)])
|
||||
assert result.exit_code != 0
|
||||
assert "exactly one of --output or --pool is required" in result.output
|
||||
|
||||
|
||||
def test_create_manifest_requires_type_with_pool(tmp_path):
|
||||
f = tmp_path / "a.parquet"
|
||||
f.touch()
|
||||
result = runner.invoke(app, ["create-manifest", "--pool", "pbwo4", str(f)])
|
||||
assert result.exit_code != 0
|
||||
assert "--type is required when --pool is given" in result.output
|
||||
|
||||
|
||||
def test_status_reports_missing_root(tmp_path):
|
||||
missing = tmp_path / "does-not-exist"
|
||||
result = runner.invoke(app, ["status", "--root", str(missing)])
|
||||
assert result.exit_code != 0
|
||||
assert "is not a directory" in result.output
|
||||
@@ -12,13 +12,17 @@ def _frame() -> pl.DataFrame:
|
||||
"track_id": [1, 1, 2, 1, 2, 3],
|
||||
"step_no": [0, 1, 0, 0, 0, 0],
|
||||
"pre_E": [100.0, 80.0, 15.0, 200.0, 20.0, 30.0],
|
||||
"pdg": [11, 11, 22, 11, 22, 22],
|
||||
"pre_dx": [0.0, 0.0, 1.0, 0.0, 1.0, 0.0],
|
||||
"pre_dy": [0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
|
||||
"pre_dz": [1.0, 1.0, 0.0, 1.0, 0.0, 0.0],
|
||||
"child_track_ids": [[2], [], [], [2, 3], [], []],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_e_sec_sums_child_first_step_energy():
|
||||
out = steps_to_parquet._add_secondary_energy(_frame())
|
||||
out = steps_to_parquet._add_secondary_attributes(_frame())
|
||||
e_sec = dict(
|
||||
zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"])
|
||||
)
|
||||
@@ -27,7 +31,7 @@ def test_e_sec_sums_child_first_step_energy():
|
||||
|
||||
|
||||
def test_e_sec_zero_when_no_children():
|
||||
out = steps_to_parquet._add_secondary_energy(_frame())
|
||||
out = steps_to_parquet._add_secondary_attributes(_frame())
|
||||
childless = out.filter(
|
||||
(pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1)
|
||||
)
|
||||
@@ -36,6 +40,6 @@ def test_e_sec_zero_when_no_children():
|
||||
|
||||
def test_e_sec_preserves_row_count_and_order():
|
||||
df = _frame()
|
||||
out = steps_to_parquet._add_secondary_energy(df)
|
||||
out = steps_to_parquet._add_secondary_attributes(df)
|
||||
assert out.height == df.height
|
||||
assert out["pre_E"].to_list() == df["pre_E"].to_list()
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from scripts import steps_to_parquet_parallel
|
||||
@@ -40,7 +41,7 @@ def test_runs_one_job_per_file_and_reports_success(tmp_path):
|
||||
fake = _write_fake_executable(tmp_path, marker_dir)
|
||||
|
||||
files = [str(tmp_path / f"shard-{i:03d}.root") for i in range(3)]
|
||||
results = run_parallel(files, jobs=4, steps_to_parquet_path=fake)
|
||||
results = run_parallel(files, jobs=4, cmd_prefix=[sys.executable, str(fake)])
|
||||
|
||||
assert {r[0] for r in results} == set(files)
|
||||
assert all(code == 0 for _, code, _, _ in results)
|
||||
@@ -53,7 +54,7 @@ def test_failures_are_reported_with_nonzero_exit_code(tmp_path):
|
||||
fake = _write_fake_executable(tmp_path, marker_dir)
|
||||
|
||||
files = [str(tmp_path / "shard-000.root"), str(tmp_path / "shard-fail.root")]
|
||||
results = run_parallel(files, jobs=4, steps_to_parquet_path=fake)
|
||||
results = run_parallel(files, jobs=4, cmd_prefix=[sys.executable, str(fake)])
|
||||
|
||||
codes = {Path(f).stem: code for f, code, _, _ in results}
|
||||
assert codes["shard-000"] == 0
|
||||
@@ -66,7 +67,7 @@ def test_jobs_caps_concurrency(tmp_path):
|
||||
fake = _write_fake_executable(tmp_path, marker_dir)
|
||||
|
||||
files = [str(tmp_path / f"shard-{i:03d}.root") for i in range(6)]
|
||||
run_parallel(files, jobs=2, steps_to_parquet_path=fake)
|
||||
run_parallel(files, jobs=2, cmd_prefix=[sys.executable, str(fake)])
|
||||
|
||||
intervals = []
|
||||
for f in files:
|
||||
@@ -83,11 +84,6 @@ def test_jobs_caps_concurrency(tmp_path):
|
||||
assert peak <= 2
|
||||
|
||||
|
||||
def test_default_jobs_is_four():
|
||||
args = steps_to_parquet_parallel.build_parser().parse_args(["dummy.root"])
|
||||
assert args.jobs == 4
|
||||
|
||||
|
||||
def test_output_for_is_passed_through_as_output_flag(tmp_path):
|
||||
marker_dir = tmp_path / "markers"
|
||||
marker_dir.mkdir()
|
||||
@@ -108,7 +104,10 @@ sys.exit(0)
|
||||
root_file = str(tmp_path / "shard-000.root")
|
||||
dest = tmp_path / "processed" / "shard-000.parquet"
|
||||
run_parallel(
|
||||
[root_file], jobs=1, steps_to_parquet_path=fake, output_for={root_file: dest}
|
||||
[root_file],
|
||||
jobs=1,
|
||||
cmd_prefix=[sys.executable, str(fake)],
|
||||
output_for={root_file: dest},
|
||||
)
|
||||
assert (marker_dir / "shard-000.txt").read_text() == str(dest)
|
||||
|
||||
@@ -125,13 +124,31 @@ def _make_dataset(tmp_path: Path, schemas: list[str] | None = None) -> Path:
|
||||
def test_resolve_destination_uses_latest_schema(tmp_path):
|
||||
root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3", "schema2"])
|
||||
dest = resolve_destination(root_file, tmp_path, schema_override=None)
|
||||
assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
|
||||
assert (
|
||||
dest
|
||||
== tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema3"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_destination_schema_override_wins(tmp_path):
|
||||
root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3"])
|
||||
dest = resolve_destination(root_file, tmp_path, schema_override="schema9")
|
||||
assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema9" / "pbwo4" / "shard-000.parquet"
|
||||
assert (
|
||||
dest
|
||||
== tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema9"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_destination_errors_without_any_schema(tmp_path):
|
||||
@@ -171,9 +188,3 @@ def test_resolve_destination_errors_on_wrong_shape(tmp_path):
|
||||
|
||||
def test_latest_schema_tag_returns_none_when_missing(tmp_path):
|
||||
assert latest_schema_tag(tmp_path / "does" / "not" / "exist") is None
|
||||
|
||||
|
||||
def test_dataset_root_and_schema_flags_default(tmp_path):
|
||||
args = steps_to_parquet_parallel.build_parser().parse_args(["dummy.root"])
|
||||
assert args.dataset_root == "/ceph/lbogner/geant_steps"
|
||||
assert args.schema is None
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
import numpy as np
|
||||
import pytest
|
||||
from giant.data.transforms import (
|
||||
energy_simplex_decode,
|
||||
energy_simplex_encode,
|
||||
inv_local_frame_rotation,
|
||||
inv_log_transform,
|
||||
local_frame_rotation,
|
||||
log_transform,
|
||||
@@ -53,6 +55,35 @@ def test_local_frame_rotation_preserves_norm():
|
||||
np.testing.assert_allclose(np.linalg.norm(result, axis=1), 1.0, atol=1e-5)
|
||||
|
||||
|
||||
def test_local_frame_rotation_rejects_near_zero_pre_dir():
|
||||
"""A degenerate (near-zero-norm) pre_dir has no well-defined frame — must
|
||||
raise instead of silently falling back to an arbitrary rotation axis."""
|
||||
pre_dir = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
post_dir = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32)
|
||||
with pytest.raises(ValueError, match="near-zero norm"):
|
||||
local_frame_rotation(pre_dir, post_dir)
|
||||
with pytest.raises(ValueError, match="near-zero norm"):
|
||||
inv_local_frame_rotation(pre_dir, post_dir)
|
||||
|
||||
|
||||
def test_local_frame_rotation_normalizes_non_unit_pre_dir():
|
||||
"""A pre_dir with float32-drift norm (not exactly 1) must still produce the
|
||||
same result as its exactly-normalized counterpart, not a skewed frame."""
|
||||
rng = np.random.default_rng(9)
|
||||
N = 50
|
||||
pre_dir_unit = rng.standard_normal((N, 3)).astype(np.float32)
|
||||
pre_dir_unit /= np.linalg.norm(pre_dir_unit, axis=1, keepdims=True)
|
||||
post_dir = rng.standard_normal((N, 3)).astype(np.float32)
|
||||
post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)
|
||||
|
||||
pre_dir_scaled = pre_dir_unit * rng.uniform(0.9, 1.1, size=(N, 1)).astype(
|
||||
np.float32
|
||||
)
|
||||
expected = local_frame_rotation(pre_dir_unit, post_dir)
|
||||
result = local_frame_rotation(pre_dir_scaled, post_dir)
|
||||
np.testing.assert_allclose(result, expected, atol=1e-4)
|
||||
|
||||
|
||||
def test_travel_direction_is_unit_norm():
|
||||
rng = np.random.default_rng(5)
|
||||
N = 50
|
||||
@@ -128,7 +159,7 @@ def test_energy_simplex_conservation():
|
||||
def test_energy_simplex_roundtrip():
|
||||
"""Encode → decode recovers energies whose lost part already sums to delta_e."""
|
||||
rng = np.random.default_rng(12)
|
||||
N = 500
|
||||
N = 500000
|
||||
pre_E = rng.uniform(1.0, 100.0, N).astype(np.float32)
|
||||
post_E = (pre_E * rng.uniform(0.0, 1.0, N)).astype(np.float32)
|
||||
delta_e = pre_E - post_E
|
||||
@@ -170,5 +201,7 @@ def test_normalizer_serialization():
|
||||
X = rng.standard_normal((50, 6)).astype(np.float32)
|
||||
norm = Normalizer().fit(X)
|
||||
norm2 = Normalizer.from_dict(norm.to_dict())
|
||||
assert norm2.mean is not None and norm.mean is not None
|
||||
assert norm2.std is not None and norm.std is not None
|
||||
np.testing.assert_allclose(norm2.mean, norm.mean)
|
||||
np.testing.assert_allclose(norm2.std, norm.std)
|
||||
|
||||
@@ -464,14 +464,20 @@ cuda = [
|
||||
{ name = "torch", version = "2.3.1+cu118", source = { registry = "https://download.pytorch.org/whl/cu118" } },
|
||||
]
|
||||
dev = [
|
||||
{ name = "awkward" },
|
||||
{ name = "ipykernel" },
|
||||
{ name = "matplotlib" },
|
||||
{ name = "polars" },
|
||||
{ name = "pytest" },
|
||||
{ name = "ruff" },
|
||||
{ name = "ty" },
|
||||
{ name = "uproot" },
|
||||
]
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "awkward", marker = "extra == 'convert'", specifier = ">=2.6,<3" },
|
||||
{ name = "giant", extras = ["convert", "analysis"], marker = "extra == 'dev'" },
|
||||
{ name = "ipykernel", marker = "extra == 'analysis'", specifier = ">=7.3.0" },
|
||||
{ name = "matplotlib", marker = "extra == 'analysis'", specifier = ">=3.8,<4" },
|
||||
{ name = "numpy", specifier = ">=1.26,<3" },
|
||||
|
||||
Reference in New Issue
Block a user