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

Fix ruff, ty, and pytest failures; apply ruff format

Browse Source 
Removes unused imports and an ambiguous variable name, narrows
Optional types before use so ty's flow analysis is satisfied, swaps
sum() over polars expressions for pl.sum_horizontal to avoid the
Literal[0] fallback type, and converts numpy bin edges to plain lists
before passing to matplotlib's hist (whose stub only accepts
Sequence[float]). Also applies ruff format across the repo, which had
drifted out of sync with the formatter.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-02 16:59:22 +02:00
parent e5bf7c51cb
commit 25718f175e
18 changed files with 469 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
analysis/export_event_observables.py
+12 -6
View File
@@ -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"),
analysis/export_photon_edep.py
+41 -5
View File
@@ -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")
analysis/export_photon_edep_by_process.py
+9 -2
View File
@@ -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)
analysis/export_photon_edep_ev.py
+8 -2
View File
@@ -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")
analysis/export_presentation_plots.py
+14 -2
View File
@@ -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()
analysis/validation.ipynb
+2 -2
View File
@@ -38,7 +38,7 @@
"\n",
"# Predict parquet produced by `giant predict --coord local --checkpoint ...`,\n",
"# carrying both pred_*/true_* columns so real vs. generated can be compared.\n",
"FILE = \"/home/lars/Programming/giant/20260629T131138.parquet\"\n"
"FILE = \"/home/lars/Programming/giant/20260629T131138.parquet\""
]
},
{
@@ -143,7 +143,7 @@
"# correlation matrices, ...), not just lazy aggregates, so they're built on\n",
"# a SampleCollection rather than the FILE path directly (see the next cell).\n",
"from giant.analysis import plot_marginals, plot_correlation_matrices, plot_pairwise\n",
"from giant.analysis import plot_direction_alignment, plot_constraint_violations\n"
"from giant.analysis import plot_direction_alignment, plot_constraint_violations"
]
},
{
giant/analysis.py
+33 -13
View File
@@ -373,7 +373,9 @@ def marginal_table(
# ---------------------------------------------------------------------------
def _kl_from_counts(real_counts: np.ndarray, gen_counts: np.ndarray, eps: float = 1e-8) -> float:
def _kl_from_counts(
real_counts: np.ndarray, gen_counts: np.ndarray, eps: float = 1e-8
) -> float:
"""KL(real || gen) from two aligned histogram bin-count arrays.
Same smoothing/normalization as `giant.validate._histogram_kl`, just
@@ -426,7 +428,9 @@ def _add_group_label(
return lf.with_columns(pl.lit("all").alias("_group"))
if group_by == "pdg":
return lf.with_columns(
(pl.lit("pdg=") + pl.col("pdg").cast(pl.Int64).cast(pl.Utf8)).alias("_group")
(pl.lit("pdg=") + pl.col("pdg").cast(pl.Int64).cast(pl.Utf8)).alias(
"_group"
)
)
if group_by == "material":
return lf.with_columns(
@@ -436,7 +440,9 @@ def _add_group_label(
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
labels = [f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})" 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]))
@@ -520,10 +526,16 @@ def _dim_hist_counts(
"""
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)
((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)
((pl.col("gen") - pl.col("lo")) / width * bins)
.floor()
.cast(pl.Int64)
.clip(0, bins - 1)
)
joined = narrow.join(lo_hi.lazy(), on="_group")
@@ -842,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")
@@ -906,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 = (
@@ -1208,8 +1224,12 @@ def compute_event_observables_pl(
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_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
)
@@ -1344,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,
@@ -1373,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,
@@ -1555,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]")
giant/cli.py
+1 -3
View File
@@ -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*:
scripts/bump_dataset_version.py
+57 -16
View File
@@ -92,7 +92,11 @@ 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*,
@@ -120,7 +124,12 @@ 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):
@@ -140,7 +149,9 @@ def plan_bump_schema(
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
@@ -196,7 +207,9 @@ def _manifest_referenced_files(pools_root: Path) -> set[Path]:
return referenced
def _referenced_root_count(raw_gen_dir: Path, processed_gen_dir: Path) -> tuple[int, int]:
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
@@ -349,13 +362,17 @@ def print_status(root: Path) -> None:
)
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)
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)
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
@@ -367,14 +384,22 @@ def print_status(root: Path) -> None:
print(_reason_line(gen_reason, indent=2))
print(
_row(
"raw", raw_size, indent=2, level="bucket",
count=raw_files, referenced=raw_referenced,
"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,
"processed",
processed_size,
indent=2,
level="bucket",
count=processed_files,
referenced=processed_referenced,
)
)
if schemas:
@@ -382,8 +407,12 @@ def print_status(root: Path) -> None:
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,
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}"))
@@ -391,7 +420,9 @@ def print_status(root: Path) -> None:
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(
_row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files)
)
print()
grand_total += kind_total
grand_files += kind_files
@@ -414,6 +445,7 @@ def print_status(root: Path) -> None:
# update-manifest
# ---------------------------------------------------------------------------
def plan_update_manifest(
manifest_path: Path,
target_schema: str | None,
@@ -489,8 +521,13 @@ def plan_update_manifest(
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")
@@ -498,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 = []
@@ -571,6 +609,7 @@ def apply_create_manifest(output_path: Path, lines: list[str]) -> None:
# CLI entry points (called from scripts/dwarf.py)
# ---------------------------------------------------------------------------
def run_status(root: str) -> None:
root_path = Path(root)
if not root_path.is_dir():
@@ -597,7 +636,9 @@ def _run_bump(
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)
new_dirs, log_line = plan_bump_schema(
root_path, kind, gen, reason, by, date, to
)
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
print("new directories:")
scripts/create_root_files.py
+64 -15
View File
@@ -61,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
@@ -111,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)]
@@ -123,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)
@@ -166,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
}
@@ -212,8 +241,19 @@ def run_make_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"
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 execute:
@@ -223,8 +263,14 @@ def run_make_root(
tmp_root = dataset_root_path / ".sim-tmp"
tmp_root.mkdir(parents=True, exist_ok=True)
results = run_all(
planned_jobs, executable, events_per_file, dataset_root_path, kind, gen,
max_workers=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()
@@ -233,7 +279,10 @@ def run_make_root(
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)
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.")
scripts/dwarf.py
+59 -22
View File
@@ -51,9 +51,7 @@ class PoolType(str, Enum):
@app.command()
def convert(
root_files: Annotated[
list[Path], typer.Argument(help="Input ROOT file(s)")
],
root_files: Annotated[list[Path], typer.Argument(help="Input ROOT file(s)")],
output: Annotated[
Optional[Path],
typer.Option(
@@ -107,11 +105,15 @@ def convert(
typer.echo("error: --jobs must be >= 1", err=True)
raise typer.Exit(1)
compression_value = "uncompressed" if compression is Compression.none else compression.value
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)
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(
@@ -149,7 +151,8 @@ def migrate(
execute: Annotated[
bool,
typer.Option(
"--execute", help="Actually move/copy files and write manifests (default: dry run)"
"--execute",
help="Actually move/copy files and write manifests (default: dry run)",
),
] = False,
copy: Annotated[
@@ -168,25 +171,40 @@ def migrate(
@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",
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)")
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)"
"--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,
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
kind=kind,
reason=reason,
by=by,
date=date,
execute=execute,
root=str(root),
to=to,
)
@@ -194,12 +212,15 @@ def bump_gen(
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",
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)")
Optional[str],
typer.Option("--date", help="Override date (default: today, ISO)"),
] = None,
to: Annotated[
Optional[str],
@@ -209,8 +230,12 @@ def bump_schema(
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,
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(
@@ -227,7 +252,9 @@ def bump_schema(
@app.command()
def status(
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
root: Annotated[
Path, typer.Option("--root", help="Dataset root")
] = _DATASET_ROOT_DEFAULT,
) -> None:
"""List existing gens/schemas per kind."""
run_status(str(root))
@@ -248,10 +275,13 @@ def update_manifest(
] = None,
gen: Annotated[
Optional[str],
typer.Option("--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"),
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)")
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."""
@@ -278,7 +308,9 @@ def create_manifest(
] = None,
type_: Annotated[
Optional[PoolType],
typer.Option("--type", help="Pool type — full, holdout, or dev (required with --pool)"),
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)")
@@ -323,7 +355,9 @@ def make_root(
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",
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,
@@ -331,7 +365,10 @@ def make_root(
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)")
bool,
typer.Option(
"--execute", help="Actually run jobs (default: dry run / print plan)"
),
] = False,
) -> None:
"""Generate new ROOT shards via a minicalosim executable."""
scripts/steps_to_parquet_parallel.py
+4 -1
View File
@@ -202,7 +202,10 @@ def run_parallel_job(
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)
raise SystemExit(1)
test-cuda.py
+1
View File
@@ -1,4 +1,5 @@
import torch
import torch.version
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
tests/test_bump_dataset_version.py
+108 -18
View File
@@ -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"]
@@ -76,7 +81,9 @@ def test_bump_schema_rejects_nonexistent_gen(tmp_path):
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")
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
@@ -93,7 +100,13 @@ 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"
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
@@ -102,7 +115,9 @@ def test_bump_schema_to_specific_tag(tmp_path):
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")
plan_bump_schema(
tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3"
)
assert False, "expected SystemExit"
except SystemExit:
pass
@@ -131,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)
@@ -138,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"
@@ -160,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"
@@ -189,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"
@@ -204,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"
@@ -220,7 +268,15 @@ def test_update_manifest_preserves_comments_and_blanks(tmp_path):
def test_update_manifest_bumps_gen(tmp_path):
parquet = tmp_path / "processed" / "steps" / "gen2" / "schema1" / "pbwo4" / "shard-000.parquet"
parquet = (
tmp_path
/ "processed"
/ "steps"
/ "gen2"
/ "schema1"
/ "pbwo4"
/ "shard-000.parquet"
)
_make_parquet(parquet)
manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -237,7 +293,15 @@ def test_update_manifest_bumps_gen(tmp_path):
def test_update_manifest_bumps_gen_and_schema(tmp_path):
parquet = tmp_path / "processed" / "steps" / "gen2" / "schema3" / "pbwo4" / "shard-000.parquet"
parquet = (
tmp_path
/ "processed"
/ "steps"
/ "gen2"
/ "schema3"
/ "pbwo4"
/ "shard-000.parquet"
)
_make_parquet(parquet)
manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -254,7 +318,15 @@ def test_update_manifest_bumps_gen_and_schema(tmp_path):
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"
@@ -274,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)
@@ -285,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)
@@ -316,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)
tests/test_cli_predict.py
+11 -3
View File
@@ -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,
)
# ---------------------------------------------------------------------------
@@ -110,7 +114,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 +131,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("/")
tests/test_create_root_files.py
+23 -8
View File
@@ -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]
)
tests/test_steps_to_parquet_parallel.py
+20 -2
View File
@@ -124,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):
tests/test_transforms.py
+2
View File
@@ -201,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)