diff --git a/CLAUDE.md b/CLAUDE.md index bdd00a5..4833588 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -101,7 +101,7 @@ Secondary energies are a **stick-breaking partition of the `e_sec` budget** from **Input is one or more `giant rollout` YAML sidecars** (`condor.py:load_rollout_yamls`): each YAML's `output`/`dataset` keys name its rollout parquet and seed file (= the reference truth); every supplied YAML must resolve to the same `dataset`, checked up front with a clear error otherwise (the premise is "N candidates vs one ground truth"). Each rollout's series name comes from a repeated `--label` CLI flag, else the YAML stem (N>1), else `"rollout"` (a single YAML). `prep` creates a **run directory** (`/analysis_runs/analysis_/` by default, `--run-dir` to override) holding `shared.json`, `run_meta.json` (`RunMeta.rollouts: list[{name,path,plot_meta}]`, insertion order = CLI order = every plot's series order), `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze submit a.yaml [b.yaml ...] --chunks N` runs `prep` (recording `N` in `run_meta.json`) then submits one HTCondor job per (plot, chunk) pair (`compute-one --id --chunk --run-dir`, polars/numpy only — no LaTeX on workers), each streaming over an `event_id`-disjoint slice (`event_id % N == chunk`) of the reference **and every rollout** and writing a small `reduced_partial/__.json`; every `PlotSpec` splits into a `compute_partial`/`finalize` pair so chunks can be summed/concatenated back per rollout (`chunkable=False` specs — the checkpoint-bound diagnostics, already bounded/subsampled — always run as a single chunk). The local `giant analyze render ` first joins every plot's chunk partials into `reduced/.json` (`merge_all`, a no-op join when `N=1`; `merge-one` does a single plot for debugging), then turns those into the styled PDF/gallery tree. `giant analyze metrics ` is a separate, unrelated entry point: training-progress plots straight from a run's `metrics.csv`. -**`prediction` family (paired truth/pred, `giant/analysis/prediction.py`):** an optional add-on to the rollout comparison, driven by `--prediction`/`--prediction-label` on `analyze prep`/`submit` (repeatable, same convention as `--label`/rollout YAMLs; series name defaults to the YAML stem for N>1 or `"prediction"` for one). Unlike a rollout (freely generated, no row-level correspondence to truth), a `giant predict` output has a matching truth row for every prediction — a paired, not distributional, comparison. `giant predict --coord global` (schema v3, `--truth` on by default) writes both `pred_*` and `true_*` physical columns plus truth/predicted secondary lists; `--coord local` is the older, always-paired 9D model-space output (`pred_{name}`/`true_{name}` for `LOCAL_TARGET_NAMES`, no secondaries — stage 2 doesn't run there). `paired_frame()` normalizes either coord into one canonical `true_`/`pred_` frame over `PAIRED_VARS` (`step_length`, `edep`, `delta_e`, `post_E`, `cos_scatter`, `cos_travel`), decoding local coord's ALR energy logits the same way `energy_simplex_decode` does. Every prediction in one run must share one `--coord` and the rollouts' `dataset` (`condor.load_prediction_yamls`). The catalog's `prediction` family (`catalog.py`, ids prefixed `pred_`) covers per-variable marginals (new `paired_hist` kind: true dashed / pred solid) and truth-vs-pred 2D histograms (new `heatmap2d` kind, with a y=x guide), residuals/relative-residuals/residual-vs-truth profiles, KS/bias/RMSE scorecards (reusing `heatmap`), `n_sec` and secondary-species confusion matrices (row-normalised `heatmap`), direction-alignment and physical-constraint-violation checks, and a pred/true correlation-matrix delta. Every spec degrades to `kind="unavailable"` when no `--prediction` was given, so a rollout-only run is unaffected. +**`prediction` family (paired truth/pred, `giant/analysis/prediction.py`):** an optional add-on to the rollout comparison, driven by `--prediction`/`--prediction-label` on `analyze prep`/`submit` (repeatable, same convention as `--label`/rollout YAMLs; series name defaults to the YAML stem for N>1 or `"prediction"` for one). Unlike a rollout (freely generated, no row-level correspondence to truth), a `giant predict` output has a matching truth row for every prediction — a paired, not distributional, comparison. `giant predict --coord global` (schema v3, `--truth` on by default) writes both `pred_*` and `true_*` physical columns plus truth/predicted secondary lists; `--coord local` is the older, always-paired 9D model-space output (`pred_{name}`/`true_{name}` for `LOCAL_TARGET_NAMES`, no secondaries — stage 2 doesn't run there). `paired_frame()` normalizes either coord into one canonical `true_`/`pred_` frame over `PAIRED_VARS` (`step_length`, `edep`, `delta_e`, `post_E`, `cos_scatter`, `cos_travel`), decoding local coord's ALR energy logits the same way `energy_simplex_decode` does. Every prediction in one run must share one `--coord` and the rollouts' `dataset` (`condor.load_prediction_yamls`). The catalog's `prediction` family (`catalog.py`, ids prefixed `pred_`) covers per-variable marginals (new `paired_hist` kind: true dashed / pred solid) and truth-vs-pred 2D histograms (new `heatmap2d` kind, with a y=x guide), residuals/relative-residuals/residual-vs-truth profiles, KS/bias/RMSE scorecards (reusing `heatmap`), `n_sec` and secondary-species confusion matrices (row-normalised `heatmap`), direction-alignment and physical-constraint-violation checks, and a pred/true correlation-matrix delta. Every spec degrades to `kind="unavailable"` when no `--prediction` was given, so a rollout-only run is unaffected. `giant predict` also writes a YAML sidecar next to the checkpoint (`cli.py:_write_prediction_ref`, mirroring `giant rollout`'s) carrying the run's provenance and timing — coord/weights/steps/batch size, row/skip/unknown-PDG counts, a `timing` block, and the checkpoint's `model_config`/`training_epoch`/`training_config` — which `--prediction` consumes the same way `--label` rollout YAMLs are consumed, surfacing those keys into each plot's gallery `metadata.yaml` (`condor.py:_PLOT_META_KEYS`). **Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower, advancing tracks breadth-first (every sweep steps all active tracks once, in `batch_size` chunks, so many tracks share each forward pass). Each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on one of the `TERM_*` reasons in `constants.py` (energy cutoff, max steps, escape, natural end, unknown pdg, max tracks); energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction. `giant/checkpoint_io.py` is the shared checkpoint → ready-to-run-models path used by both `predict` and `rollout`. diff --git a/giant/analysis/condor.py b/giant/analysis/condor.py index 2bd85b9..6437c40 100644 --- a/giant/analysis/condor.py +++ b/giant/analysis/condor.py @@ -61,7 +61,9 @@ from giant.analysis.reduced import Partial from giant.analysis.runtime_estimate import estimate_runtime_s from giant.analysis.sources import RolloutSpec, Side, open_side -# Keys copied verbatim from a rollout YAML into each plot's gallery metadata. +# Keys copied verbatim from a rollout or prediction YAML into each plot's +# gallery metadata. Rollout-only and predict-only keys both live here — +# `_plot_meta` copies only whichever of these are present in a given YAML. _PLOT_META_KEYS = ( "prediction_id", "checkpoint", @@ -85,10 +87,20 @@ _PLOT_META_KEYS = ( "termination_reason_counts", "timing", "model_config", + "config_overrides", "training_epoch", "best_val_loss", "training_config", "training_meta", + # giant predict only (giant/cli.py's predict command). + "coord", + "has_truth", + "schema_version", + "n_input_rows", + "n_files", + "n_skipped_rows", + "unknown_pdg_counts", + "batch_size_auto", # Diagnostic — only present when giant rollout ran under # stage2_model.particle_type.target="embedding" (see giant/cli.py's # rollout command and giant.rollout.L1DistCollector); absent otherwise, diff --git a/giant/cli.py b/giant/cli.py index 916d728..66ed0f9 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -187,8 +187,14 @@ def _write_prediction_ref( out: Path, dataset_path: Path, comment: str | None = None, + extra: dict | None = None, ) -> Path: - """Write a YAML sidecar in the checkpoint directory and return its path.""" + """Write a YAML sidecar in the checkpoint directory and return its path. + + ``extra`` is merged in after the base fields (e.g. `giant rollout`'s + provenance/timing block, or `giant predict`'s) — callers that don't pass + it get exactly today's thin sidecar. + """ import yaml ref = { @@ -201,6 +207,8 @@ def _write_prediction_ref( } if comment is not None: ref["comment"] = comment + if extra: + ref.update(extra) 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,6 +253,41 @@ def _build_rollout_timing( } +def _build_predict_timing( + *, + setup_s: float, + predict_s: float, + write_s: float, + n_rows: int, + device: str, + torch_threads: int, +) -> dict: + """Assemble ``giant predict``'s ``timing`` sidecar block. + + Keys are deliberately compatible with ``_build_rollout_timing``'s (same + names for the quantities both commands have) so a gallery's ``timing`` + metadata renders the same way whether the series came from a rollout or + a prediction. There's no ``n_physical_rows``/``ms_per_event`` here: + unlike a rollout, `giant predict` never emits synthetic termination rows + (one output row per input step) and doesn't work in whole showers/events + — so ``us_per_step`` is already directly comparable to a rollout's and to + ``giant.analysis.geant4_reference``'s per-step Geant4 measurement. + """ + sample_s = predict_s - write_s + return { + "setup_s": setup_s, + "predict_s": predict_s, + "write_s": write_s, + "sample_s": sample_s, + "n_rows": n_rows, + "us_per_step": (sample_s / n_rows * 1e6) if n_rows else None, + "write_us_per_step": (write_s / n_rows * 1e6) if n_rows else None, + "rows_per_s": (n_rows / predict_s) if predict_s else None, + "device": device, + "torch_threads": torch_threads, + } + + @app.callback() def _main() -> None: """GIANT — Geant4 step-function surrogate.""" @@ -1085,6 +1128,8 @@ def predict( ] = None, ) -> None: """Run trained model on a parquet file and save predictions.""" + import time + import numpy as np import pyarrow as pa import pyarrow.parquet as pq @@ -1104,6 +1149,8 @@ def predict( from giant.rollout import decode_secondary_identity from giant.sample import resolve_n_sec, sample_stage1, sample_stage2 + _t_setup_start = time.perf_counter() + batch_size_auto = False batch_size_value: int | None = None if batch_size.strip().lower() == "auto": @@ -1175,6 +1222,10 @@ def predict( unknown_pdg_counts: Counter[int] = Counter() total_rows = sum(pq.ParquetFile(path).metadata.num_rows for path in files) + training_cfg = gconfig.load_checkpoint_config(checkpoint) + _write_s = 0.0 + _setup_s = time.perf_counter() - _t_setup_start + # --coord local always needs full row-groups (it's paired against the 9D # target); --coord global only needs them when --truth is requested — # otherwise the cheaper conditioning-only read is used. @@ -1196,7 +1247,7 @@ def predict( return {k: np.concatenate([a[k], b[k]], axis=0) for k in a} def _process(piece: dict[str, np.ndarray]) -> None: - nonlocal writer, total + nonlocal writer, total, _write_s if coord == Coord.local: feats = build_features( @@ -1386,15 +1437,18 @@ def predict( } ) + _t0 = time.perf_counter() if writer is None: writer = pq.ParquetWriter(out, table.schema) writer.write_table(table) + _write_s += time.perf_counter() - _t0 total += len(piece["event_id"]) # Buffer rows across row-group boundaries so the inference batch size # isn't capped by however the source file happens to be chunked. buffer: dict[str, np.ndarray] | None = None + _t_predict_start = time.perf_counter() bar = tqdm(total=total_rows, desc="predict", unit="row", dynamic_ncols=True) for i, path in enumerate(files): for chunk in chunk_iter(path, offset=event_id_offset(i)): @@ -1422,9 +1476,56 @@ def predict( bar.close() if writer is not None: writer.close() + _predict_s = time.perf_counter() - _t_predict_start - ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path, comment) + timing = _build_predict_timing( + setup_s=_setup_s, + predict_s=_predict_s, + write_s=_write_s, + n_rows=total, + device=str(_device), + torch_threads=torch.get_num_threads(), + ) + + ref_path = _write_prediction_ref( + checkpoint, + pred_uuid, + out, + dataset_path, + comment, + extra={ + "coord": coord.value, + "has_truth": coord == Coord.local or write_truth, + "schema_version": PREDICT_SCHEMA_VERSION, + "steps": steps, + "weights": weights.value, + "device": str(_device), + "batch_size": bs, + "batch_size_auto": batch_size_auto, + "n_input_rows": total_rows, + "n_files": len(files), + "n_rows": total, + "n_skipped_rows": skipped, + "unknown_pdg_counts": {str(pdg): count for pdg, count in unknown_pdg_counts.items()}, + "timing": timing, + # Full architecture spec baked into the checkpoint — see the + # matching comment in `rollout`. + "model_config": dict(ctx.model_config), + "config_overrides": dict(ctx.config_overrides), + "training_epoch": ctx.epoch, + "best_val_loss": ctx.best_val_loss, + # [train]/[meta] from the sibling config.toml (giant.config.save_config) + # — empty dicts if the checkpoint has no config.toml next to it. + "training_config": dict(training_cfg.get("train", {})), + "training_meta": dict(training_cfg.get("meta", {})), + }, + ) typer.echo(f"reference: {ref_path}") + if timing["us_per_step"] is not None: + typer.echo( + f"timing: {_predict_s:.1f}s total ({timing['sample_s']:.1f}s sample + {_write_s:.1f}s write), " + f"{timing['us_per_step']:.1f} us/step over {total:,} step(s)" + ) if skipped: codes = ", ".join(f"{pdg} ({count})" for pdg, count in sorted(unknown_pdg_counts.items())) @@ -1556,7 +1657,6 @@ def rollout( import pyarrow as pa import pyarrow.parquet as pq import torch - import yaml from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference from giant.data.loader import find_parquet_files @@ -1685,10 +1785,12 @@ def rollout( l1_summary = l1_dist_collector.summary() - ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path) - ref = yaml.safe_load(ref_path.read_text()) - ref.update( - { + ref_path = _write_prediction_ref( + checkpoint, + pred_uuid, + out, + dataset_path, + extra={ "kind": "rollout", "geometry_oracle": str(geometry.resolve()), "energy_cutoff": energy_cutoff, @@ -1725,9 +1827,8 @@ def rollout( # — empty dicts if the checkpoint has no config.toml next to it. "training_config": dict(training_cfg.get("train", {})), "training_meta": dict(training_cfg.get("meta", {})), - } + }, ) - ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False)) typer.echo(f"wrote {summary['n_rows']:,} step rows → {out}") typer.echo(f"terminations: {summary['termination_reason_counts']}") diff --git a/tests/test_cli_predict.py b/tests/test_cli_predict.py index 7459e1c..16c1f5b 100644 --- a/tests/test_cli_predict.py +++ b/tests/test_cli_predict.py @@ -6,6 +6,7 @@ from typer.testing import CliRunner from giant.cli import ( _CEPH_PREDICTIONS, + _build_predict_timing, _resolve_prediction_output, _write_prediction_ref, app, @@ -145,6 +146,46 @@ def test_ref_timestamp_is_iso_format(tmp_path): assert ts.tzinfo is not None +def test_ref_yaml_merges_extra_after_base_fields(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, + extra={"coord": "global", "n_rows": 42, "timing": {"setup_s": 1.0}}, + ) + data = yaml.safe_load(ref_path.read_text()) + + # Base fields untouched, extras layered on top. + assert data["kind"] == "prediction" + assert data["prediction_id"] == pred_uuid + assert data["coord"] == "global" + assert data["n_rows"] == 42 + assert data["timing"] == {"setup_s": 1.0} + + +def test_ref_yaml_without_extra_matches_today(tmp_path): + ckpt_dir = tmp_path / "checkpoints" + ckpt_dir.mkdir() + checkpoint = ckpt_dir / "best.pt" + checkpoint.touch() + + pred_uuid = str(uuid.uuid4()) + ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d") + data = yaml.safe_load(ref_path.read_text()) + + assert set(data) == {"kind", "prediction_id", "output", "dataset", "checkpoint", "timestamp"} + + def test_ref_checkpoint_path_is_absolute(tmp_path): ckpt_dir = tmp_path / "checkpoints" ckpt_dir.mkdir() @@ -232,13 +273,52 @@ def test_predict_truth_metadata_key_exists(): assert PREDICT_TRUTH_METADATA_KEY == "giant.predict.has_truth" +# --------------------------------------------------------------------------- +# _build_predict_timing +# --------------------------------------------------------------------------- + + +def test_build_predict_timing_computes_per_step_cost(): + timing = _build_predict_timing( + setup_s=1.0, + predict_s=10.0, + write_s=2.0, + n_rows=100, + device="cpu", + torch_threads=4, + ) + assert timing["n_rows"] == 100 + assert timing["sample_s"] == 8.0 # predict_s - write_s + assert timing["us_per_step"] == 8.0 / 100 * 1e6 + assert timing["write_us_per_step"] == 2.0 / 100 * 1e6 + assert timing["rows_per_s"] == 10.0 + assert timing["device"] == "cpu" and timing["torch_threads"] == 4 + + +def test_build_predict_timing_handles_zero_rows(): + timing = _build_predict_timing( + setup_s=1.0, + predict_s=0.0, + write_s=0.0, + n_rows=0, + device="cpu", + torch_threads=1, + ) + assert timing["us_per_step"] is None + assert timing["write_us_per_step"] is None + assert timing["rows_per_s"] is None + + def test_predict_has_truth_flag_default_on(): # Inspecting rendered --help text is brittle across terminal # widths/color settings (wraps or re-colors mid-flag); go straight to # the underlying click command's registered option instead. - import typer + from typing import cast - predict_cmd = typer.main.get_command(app).commands["predict"] + import typer + from click import Group + + predict_cmd = cast(Group, typer.main.get_command(app)).commands["predict"] truth_param = next(p for p in predict_cmd.params if p.name == "truth") assert truth_param.opts == ["--truth"] assert truth_param.secondary_opts == ["--no-truth"] diff --git a/tests/test_condor.py b/tests/test_condor.py index 2cbdc87..8ec8092 100644 --- a/tests/test_condor.py +++ b/tests/test_condor.py @@ -256,6 +256,46 @@ def test_prep_with_prediction_writes_run_meta(tmp_path: Path): assert partial.data["available"] +def test_prep_forwards_predict_only_metadata_keys(tmp_path: Path): + """A rich `giant predict` sidecar's provenance/timing keys reach + run_meta.json's plot_meta, same as a rollout's do — a thin legacy + sidecar (no such keys) still loads fine (see _write_prediction_yaml).""" + rollout_yaml = _write_inputs(tmp_path) + reference = load_rollout_yaml(rollout_yaml)["dataset"] + pred = tmp_path / "pred_rich.parquet" + _write_prediction(pred, coord="global") + yaml_path = tmp_path / "pred_rich.yaml" + yaml_path.write_text( + yaml.safe_dump( + { + "kind": "prediction", + "prediction_id": "richpred12", + "output": str(pred), + "dataset": str(reference), + "checkpoint": "/ckpt/rich.pt", + "coord": "global", + "has_truth": True, + "schema_version": "3", + "n_input_rows": 1000, + "n_files": 1, + "n_skipped_rows": 3, + "unknown_pdg_counts": {"999999": 3}, + "batch_size_auto": False, + "timing": {"us_per_step": 12.5}, + } + ) + ) + run_dir = _prep([rollout_yaml], prediction_yamls=[yaml_path]) + meta = RunMeta.load(run_dir / "run_meta.json") + plot_meta = meta.predictions[0]["plot_meta"] + assert plot_meta["coord"] == "global" + assert plot_meta["has_truth"] is True + assert plot_meta["n_input_rows"] == 1000 + assert plot_meta["n_skipped_rows"] == 3 + assert plot_meta["unknown_pdg_counts"] == {"999999": 3} + assert plot_meta["timing"] == {"us_per_step": 12.5} + + def test_derive_run_dir_next_to_rollout(): y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"} assert derive_run_dir([y]) == Path("/data/analysis_abcd1234")