feat(predict): enrich YAML sidecar with provenance and timing

`giant predict`'s sidecar previously stopped at kind/prediction_id/
output/dataset/checkpoint/timestamp, unlike `giant rollout`'s, which
carries full run provenance (model_config, training_epoch,
training_config, timing, ...) that flows into analysis gallery
metadata. `analyze --prediction` consumed the same thin sidecar, so a
prediction series in an analysis run was nearly unlabeled compared to
its rollout counterparts.

- `_write_prediction_ref` takes an `extra: dict | None` merged into
  the sidecar; `giant rollout` now uses it instead of a
  load/update/rewrite round trip (identical output).
- New `_build_predict_timing`, key-compatible with
  `_build_rollout_timing`, from timers now wrapping predict's setup/
  sample/write phases.
- `giant predict` writes coord, has_truth, schema_version, steps,
  weights, device, batch_size(+auto), row/skip/unknown-pdg counts,
  timing, and the checkpoint's model_config/config_overrides/
  training_epoch/best_val_loss/training_config/training_meta.
- `giant/analysis/condor.py`'s `_PLOT_META_KEYS` forwards the new
  predict-only keys (plus rollout's previously-unforwarded
  config_overrides) into each plot's gallery metadata.yaml.
- Fixes a `ty` regression from the prior commit in
  tests/test_cli_predict.py (Command has no static `.commands`).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PpxE9nij3ujg9XcuzvQ26q

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** (`<cwd>/analysis_runs/analysis_<id>/` 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/<id>__<chunk>.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 <run_dir>` first joins every plot's chunk partials into `reduced/<id>.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 <train_run_dir>` 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_<var>`/`pred_<var>` 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_<var>`/`pred_<var>` 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`.
 

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,

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']}")

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"]

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")