feat(analyze): add paired truth/pred plots from giant predict

Adds a `prediction` plot family to `giant analyze`, alongside the existing
rollout-vs-reference comparison, and extends `giant predict` to make it
possible:

- `giant predict --coord global` gains schema v3 (`--truth/--no-truth`,
  default on): writes true_* physical columns and true secondary lists
  alongside the predictions, so the output is fully paired.
- New `giant/analysis/prediction.py` builds one canonical true/pred frame
  (`paired_frame`) from either predict coord mode.
- `catalog.py` gains 35 `pred_*` specs: marginals, 2D truth-vs-pred scatter
  (new `heatmap2d` kind), residuals/relative-residuals/calibration profiles,
  KS/bias/RMSE scorecards, n_sec + secondary-species confusion matrices,
  direction-alignment and constraint-violation checks, and a correlation
  delta. Two new Reduced kinds (`paired_hist`, `heatmap2d`) get renderers.
  Every spec degrades to kind="unavailable" with no --prediction given.
- `condor.py`/`cli.py`: `--prediction`/`--prediction-label` on
  `analyze prep`/`submit`, threaded through RunMeta and every compute job.

Full test suite (1162 tests), ruff, and ty all pass.

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

tests/test_analysis_prediction.py

@@ -0,0 +1,278 @@
+"""Tests for giant.analysis.prediction (paired truth/pred frames for `giant predict`
+output) and the `prediction` family of catalog specs."""
+
+from __future__ import annotations
+
+import numpy as np
+import polars as pl
+import pytest
+
+from giant.analysis.catalog import Bundle, get_spec
+from giant.analysis.context import Context, build_context
+from giant.analysis.prediction import (
+    PAIRED_SCALARS,
+    PredictionSpec,
+    open_prediction,
+    paired_frame,
+    paired_secondaries,
+    prediction_secondaries,
+)
+from giant.analysis.reduce import hist2d
+from giant.analysis.sources import RolloutSpec
+from tests.test_analysis_reduce import _reference_frame, _rollout_frame
+
+
+def _global_prediction_frame() -> pl.LazyFrame:
+    """A `--coord global --truth` predict parquet, as a LazyFrame (schema per
+    `giant.cli.predict`'s global-coord table, `giant/cli.py:1310-1379`)."""
+    return pl.DataFrame(
+        {
+            "event_id": [1, 1, 2],
+            "pdg": [11, 11, 22],
+            "pre_x": [0.0, 0.0, 0.0],
+            "pre_y": [0.0, 0.0, 0.0],
+            "pre_z": [0.0, 1.0, 0.0],
+            "pre_E": [100.0, 60.0, 50.0],
+            "pre_dx": [0.0, 0.0, 0.0],
+            "pre_dy": [0.0, 0.0, 0.0],
+            "pre_dz": [1.0, 1.0, 1.0],
+            "material": ["G4_PbWO4", "G4_PbWO4", "G4_Pb"],
+            "layer_id": [0, 1, 0],
+            "n_sec": [1, 0, 2],
+            "n_sec_pred": [1, 0, 1],
+            # predicted (unprefixed) values
+            "step_length": [1.2, 0.9, 1.1],
+            "delta_e": [42.0, 29.0, 31.0],
+            "edep": [35.0, 29.0, 25.0],
+            "post_dx": [0.0, 0.0, 0.0],
+            "post_dy": [0.0, 0.0, 0.0],
+            "post_dz": [1.0, 1.0, 1.0],
+            "post_x": [0.0, 0.0, 0.0],
+            "post_y": [0.0, 0.0, 0.0],
+            "post_z": [1.2, 1.9, 1.1],
+            "sec_pdg_list": [[22], [], [22]],
+            "sec_E_list": [[5.0], [], [4.0]],
+            "sec_dx_list": [[0.0], [], [0.0]],
+            "sec_dy_list": [[0.0], [], [0.0]],
+            "sec_dz_list": [[1.0], [], [1.0]],
+            # truth
+            "true_step_length": [1.0, 1.0, 1.0],
+            "true_delta_e": [40.0, 30.0, 30.0],
+            "true_edep": [40.0, 30.0, 20.0],
+            "true_post_E": [60.0, 30.0, 20.0],
+            "true_post_dx": [0.0, 0.0, 0.0],
+            "true_post_dy": [0.0, 0.0, 0.0],
+            "true_post_dz": [1.0, 1.0, 1.0],
+            "true_post_x": [0.0, 0.0, 0.0],
+            "true_post_y": [0.0, 0.0, 0.0],
+            "true_post_z": [1.0, 2.0, 1.0],
+            "true_e_sec": [0.0, 0.0, 10.0],
+            "process": ["compt", "phot", "compt"],
+            "true_sec_pdg_list": [[22], [], [22, 11]],
+            "true_sec_E_list": [[6.0], [], [7.0, 3.0]],
+            "true_sec_dx_list": [[0.0], [], [0.0, 1.0]],
+            "true_sec_dy_list": [[0.0], [], [0.0, 0.0]],
+            "true_sec_dz_list": [[1.0], [], [1.0, 0.0]],
+        }
+    ).lazy()
+
+
+def _local_prediction_frame() -> pl.LazyFrame:
+    """A `--coord local` predict parquet — always paired, never has secondaries."""
+    return pl.DataFrame(
+        {
+            "event_id": [1, 2],
+            "pdg": [11, 22],
+            "pre_x": [0.0, 0.0],
+            "pre_y": [0.0, 0.0],
+            "pre_z": [0.0, 0.0],
+            "pre_E": [100.0, 50.0],
+            "pre_dx": [0.0, 0.0],
+            "pre_dy": [0.0, 0.0],
+            "pre_dz": [1.0, 1.0],
+            "material": ["G4_PbWO4", "G4_Pb"],
+            "layer_id": [0, 0],
+            "n_sec": [1, 0],
+            # ALR logits: [edep_logit, sec_logit] -> softmax([z1,z2,0]) * pre_E
+            "pred_log_step_length": [np.log(1.2 + 1e-6), np.log(0.9 + 1e-6)],
+            "pred_edep_logit": [1.0, 0.5],
+            "pred_sec_logit": [0.0, -1.0],
+            "pred_post_dx": [0.0, 0.0],
+            "pred_post_dy": [0.0, 0.0],
+            "pred_post_dz": [1.0, 1.0],
+            "pred_travel_dx": [0.0, 0.0],
+            "pred_travel_dy": [0.0, 0.0],
+            "pred_travel_dz": [1.0, 1.0],
+            "true_log_step_length": [np.log(1.0 + 1e-6), np.log(1.0 + 1e-6)],
+            "true_edep_logit": [0.8, 0.6],
+            "true_sec_logit": [0.2, -2.0],
+            "true_post_dx": [0.0, 0.0],
+            "true_post_dy": [0.0, 0.0],
+            "true_post_dz": [1.0, 1.0],
+            "true_travel_dx": [0.0, 0.0],
+            "true_travel_dy": [0.0, 0.0],
+            "true_travel_dz": [1.0, 1.0],
+        }
+    ).lazy()
+
+
+def test_open_prediction_detects_coord_and_truth():
+    g = open_prediction(_global_prediction_frame())
+    assert g.coord == "global" and g.has_truth
+
+    loc = open_prediction(_local_prediction_frame())
+    assert loc.coord == "local" and loc.has_truth
+
+
+def test_paired_frame_global_matches_source_columns():
+    lf = _global_prediction_frame()
+    p = paired_frame(lf, "global", has_truth=True).collect()
+    assert p["pred_step_length"].to_list() == [1.2, 0.9, 1.1]
+    assert p["true_step_length"].to_list() == [1.0, 1.0, 1.0]
+    assert p["pred_edep"].to_list() == [35.0, 29.0, 25.0]
+    assert p["true_edep"].to_list() == [40.0, 30.0, 20.0]
+    # post_E isn't written directly for the prediction (energy conservation:
+    # pre_E - delta_e); truth carries it verbatim.
+    assert p["pred_post_E"].to_list() == pytest.approx([100.0 - 42.0, 60.0 - 29.0, 50.0 - 31.0])
+    assert p["true_post_E"].to_list() == [60.0, 30.0, 20.0]
+    # cos_scatter: pre_dir . post_dir, both (0,0,1) here -> 1.0
+    assert p["pred_cos_scatter"].to_list() == pytest.approx([1.0, 1.0, 1.0])
+    assert p["true_cos_scatter"].to_list() == pytest.approx([1.0, 1.0, 1.0])
+
+
+def test_paired_frame_local_decodes_energy_simplex():
+    lf = _local_prediction_frame()
+    p = paired_frame(lf, "local", has_truth=True).collect()
+    # softmax([1.0, 0.0, 0.0]) * 100 for row 0's pred edep
+    z = np.exp([1.0, 0.0, 0.0])
+    expected_edep_0 = (z[0] / z.sum()) * 100.0
+    assert p["pred_edep"][0] == pytest.approx(expected_edep_0)
+    assert p["pred_step_length"][0] == pytest.approx(1.2, abs=1e-4)
+    # local coord never has a meaningful cos_travel (no reconstructed post_pos)
+    assert "cos_travel" not in [c.rsplit("_", 1)[-1] for c in ["pred_cos_travel"] if c in p.columns] or True
+    assert "pred_cos_travel" not in p.columns
+
+
+def test_prediction_secondaries_and_pairing():
+    lf = _global_prediction_frame()
+    true_sec = prediction_secondaries(lf, "true").collect()
+    pred_sec = prediction_secondaries(lf, "pred").collect()
+    assert true_sec["pdg"].to_list() == [22, 22, 11]
+    assert pred_sec["pdg"].to_list() == [22, 22]
+
+    pairs = paired_secondaries(lf).collect()
+    # event 1: 1 true, 1 pred -> paired (22, 22); event 2: 2 true, 1 pred -> paired rank0 only (22, 22)
+    assert pairs["true_pdg"].to_list() == [22, 22]
+    assert pairs["pred_pdg"].to_list() == [22, 22]
+
+
+def test_hist2d_basic():
+    lf = pl.DataFrame({"x": [0.1, 0.5, 0.9, 0.5], "y": [0.1, 0.9, 0.9, 0.1]}).lazy()
+    edges = np.linspace(0.0, 1.0, 3)  # 2 bins: [0,0.5), [0.5,1]
+    mat = hist2d(lf, pl.col("x"), pl.col("y"), edges, edges)
+    assert mat.sum() == 4
+    assert mat.shape == (2, 2)
+
+
+def _ctx_with_predictions(n_marginal_bins: int = 10) -> Context:
+    return build_context(
+        [RolloutSpec("rollout", _rollout_frame())],
+        _reference_frame(),
+        predictions=[PredictionSpec("pred", _global_prediction_frame())],
+        n_energy_bins=2,
+        n_marginal_bins=n_marginal_bins,
+        top_k_pdg=3,
+        sample_rows=1000,
+    )
+
+
+def test_build_context_resolves_prediction_ranges():
+    ctx = _ctx_with_predictions()
+    assert "edep" in ctx.pred_var_ranges
+    assert "edep" in ctx.pred_residual_ranges
+    assert ctx.pred_top_sec_pdgs  # secondaries present in the fixture
+
+
+def test_prediction_specs_compute_valid_reduced():
+    ctx = _ctx_with_predictions()
+    bundle = Bundle.open(
+        [RolloutSpec("rollout", _rollout_frame())],
+        _reference_frame(),
+        ctx,
+        predictions=[PredictionSpec("pred", _global_prediction_frame())],
+    )
+    for spec_id in (
+        "pred_marginal_edep",
+        "pred_scatter_edep",
+        "pred_residual_edep",
+        "pred_relative_residual_edep",
+        "pred_residual_profile_edep",
+        "pred_ks_summary",
+        "pred_bias_summary",
+        "pred_rmse_summary",
+        "pred_n_sec_confusion",
+        "pred_sec_species_confusion",
+        "pred_dir_alignment_post",
+        "pred_dir_alignment_travel",
+        "pred_constraint_violations",
+        "pred_correlation_delta",
+    ):
+        spec = get_spec(spec_id)
+        r = spec.finalize([spec.compute_partial(bundle)], ctx)
+        assert r.id == spec_id
+        assert r.kind != "unavailable", f"{spec_id} unexpectedly unavailable"
+        assert "pred" in r.payload["series"]
+
+
+def test_prediction_specs_unavailable_without_predictions():
+    ctx = _ctx_with_predictions()
+    bundle = Bundle.open([RolloutSpec("rollout", _rollout_frame())], _reference_frame(), ctx)
+    for spec_id in ("pred_marginal_edep", "pred_scatter_edep", "pred_n_sec_confusion", "pred_ks_summary"):
+        spec = get_spec(spec_id)
+        r = spec.finalize([spec.compute_partial(bundle)], ctx)
+        assert r.kind == "unavailable"
+        assert r.payload["note"]
+
+
+@pytest.mark.parametrize(
+    "spec_id",
+    ["pred_marginal_edep", "pred_scatter_edep", "pred_n_sec_confusion", "pred_ks_summary", "pred_correlation_delta"],
+)
+def test_prediction_chunked_matches_unchunked(spec_id: str):
+    ctx = _ctx_with_predictions()
+    specs = [RolloutSpec("rollout", _rollout_frame())]
+    preds = [PredictionSpec("pred", _global_prediction_frame())]
+    spec = get_spec(spec_id)
+
+    unchunked_bundle = Bundle.open(specs, _reference_frame(), ctx, predictions=preds)
+    unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], ctx)
+
+    n_chunks = 2
+    parts = [
+        spec.compute_partial(Bundle.open(specs, _reference_frame(), ctx, chunk=(k, n_chunks), predictions=preds))
+        for k in range(n_chunks)
+    ]
+    chunked = spec.finalize(parts, ctx)
+
+    assert chunked.kind == unchunked.kind
+    _assert_close(unchunked.payload, chunked.payload)
+
+
+def _assert_close(a, b) -> None:
+    """Recursively compare two JSON-shaped payloads (float-tolerant)."""
+    if isinstance(a, dict):
+        assert set(a) == set(b)
+        for k in a:
+            _assert_close(a[k], b[k])
+    elif isinstance(a, list):
+        assert len(a) == len(b)
+        for x, y in zip(a, b):
+            _assert_close(x, y)
+    elif isinstance(a, float):
+        assert np.isclose(a, b, atol=1e-9) or (np.isnan(a) and np.isnan(b))
+    else:
+        assert a == b
+
+
+def test_paired_scalars_are_subset_of_all_vars():
+    assert set(PAIRED_SCALARS) <= {"step_length", "edep", "delta_e", "post_E"}

tests/test_cli_predict.py

@@ -103,6 +103,7 @@ def test_ref_yaml_contains_expected_fields(tmp_path):
     ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset)
     data = yaml.safe_load(ref_path.read_text())
 
+    assert data["kind"] == "prediction"
     assert data["prediction_id"] == pred_uuid
     assert data["output"] == str(out)
     assert data["dataset"] == str(dataset)
@@ -212,3 +213,26 @@ def test_predict_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
 
     assert result.exit_code == 1
     assert "not an inference-safe override" in result.output
+
+
+# ---------------------------------------------------------------------------
+# schema v3 constants (truth-column tagging)
+# ---------------------------------------------------------------------------
+
+
+def test_predict_schema_version_is_v3():
+    from giant.constants import PREDICT_SCHEMA_VERSION
+
+    assert PREDICT_SCHEMA_VERSION == "3"
+
+
+def test_predict_truth_metadata_key_exists():
+    from giant.constants import PREDICT_TRUTH_METADATA_KEY
+
+    assert PREDICT_TRUTH_METADATA_KEY == "giant.predict.has_truth"
+
+
+def test_predict_has_truth_flag_default_on():
+    result = runner.invoke(app, ["predict", "--help"])
+    assert "--truth" in result.output
+    assert "--no-truth" in result.output

tests/test_condor.py

@@ -16,6 +16,8 @@ from giant.analysis import (
     compute_one,
     compute_reduced,
     derive_run_dir,
+    load_prediction_yaml,
+    load_prediction_yamls,
     load_rollout_yaml,
     load_rollout_yamls,
     merge_one,
@@ -25,7 +27,8 @@ from giant.analysis import (
 from giant.analysis.catalog import get_spec
 from giant.analysis.condor import Context
 from giant.analysis.reduced import Partial, Reduced
-from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE
+from giant.constants import PREDICT_COORD_METADATA_KEY, PREDICT_TRUTH_METADATA_KEY, ROLLOUT_COORD_VALUE
+from tests.test_analysis_prediction import _global_prediction_frame
 from tests.test_analysis_reduce import _reference_frame, _rollout_frame
 
 
@@ -86,6 +89,30 @@ def _write_two_inputs(tmp_path: Path) -> tuple[Path, Path]:
     return paths[0], paths[1]
 
 
+def _write_prediction(path: Path, coord: str = "global") -> None:
+    tbl = _global_prediction_frame().collect().to_arrow()
+    tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: coord, PREDICT_TRUTH_METADATA_KEY: "1"})
+    pq.write_table(tbl, path)
+
+
+def _write_prediction_yaml(tmp_path: Path, reference: Path, tag: str = "p", coord: str = "global") -> Path:
+    pred = tmp_path / f"pred_{tag}.parquet"
+    _write_prediction(pred, coord=coord)
+    yaml_path = tmp_path / f"pred_{tag}.yaml"
+    yaml_path.write_text(
+        yaml.safe_dump(
+            {
+                "kind": "prediction",
+                "prediction_id": f"{tag}pred1234",
+                "output": str(pred),
+                "dataset": str(reference),
+                "checkpoint": f"/ckpt/{tag}.pt",
+            }
+        )
+    )
+    return yaml_path
+
+
 def _fake_venv(repo_dir: Path) -> None:
     """Stand in for a `uv sync`'d venv: write_submit checks `.venv/bin/giant` exists."""
     giant = repo_dir / ".venv" / "bin" / "giant"
@@ -94,13 +121,14 @@ def _fake_venv(repo_dir: Path) -> None:
     giant.chmod(0o755)
 
 
-def _prep(rollout_yamls, run_dir: str | Path | None = None, chunks: int = 1, labels=None) -> Path:
+def _prep(rollout_yamls, run_dir: str | Path | None = None, chunks: int = 1, labels=None, prediction_yamls=()) -> Path:
     """``prep`` with small test-sized context bins/sampling."""
     return prep(
         rollout_yamls,
         run_dir,
         n_chunks=chunks,
         labels=labels,
+        prediction_yamls=prediction_yamls,
         n_energy_bins=2,
         n_marginal_bins=8,
         top_k_pdg=3,
@@ -161,6 +189,73 @@ def test_load_rollout_yamls_rejects_mismatched_reference(tmp_path: Path):
         load_rollout_yamls([a, c])
 
 
+def test_load_prediction_yaml_requires_paths(tmp_path: Path):
+    bad = tmp_path / "bad.yaml"
+    bad.write_text(yaml.safe_dump({"output": "x.parquet"}))  # no dataset
+    with pytest.raises(ValueError):
+        load_prediction_yaml(bad)
+
+
+def test_load_prediction_yaml_rejects_rollout_kind(tmp_path: Path):
+    y = tmp_path / "r.yaml"
+    y.write_text(yaml.safe_dump({"output": "x.parquet", "dataset": "d.parquet", "kind": "rollout"}))
+    with pytest.raises(ValueError, match="kind"):
+        load_prediction_yaml(y)
+
+
+def test_load_prediction_yamls_single_defaults_to_prediction_name(tmp_path: Path):
+    reference = tmp_path / "reference.parquet"
+    _reference_frame().collect().write_parquet(reference)
+    y = _write_prediction_yaml(tmp_path, reference)
+    loaded = load_prediction_yamls([y], str(reference))
+    assert [lp.name for lp in loaded] == ["prediction"]
+    assert loaded[0].coord == "global"
+
+
+def test_load_prediction_yamls_multi_defaults_to_stem_and_labels(tmp_path: Path):
+    reference = tmp_path / "reference.parquet"
+    _reference_frame().collect().write_parquet(reference)
+    a = _write_prediction_yaml(tmp_path, reference, tag="a")
+    b = _write_prediction_yaml(tmp_path, reference, tag="b")
+    loaded = load_prediction_yamls([a, b], str(reference))
+    assert [lp.name for lp in loaded] == ["pred_a", "pred_b"]
+    loaded = load_prediction_yamls([a, b], str(reference), labels=["ep20", "ep50"])
+    assert [lp.name for lp in loaded] == ["ep20", "ep50"]
+
+
+def test_load_prediction_yamls_rejects_mismatched_reference(tmp_path: Path):
+    reference = tmp_path / "reference.parquet"
+    _reference_frame().collect().write_parquet(reference)
+    other_ref = tmp_path / "other_reference.parquet"
+    _reference_frame().collect().write_parquet(other_ref)
+    y = _write_prediction_yaml(tmp_path, other_ref)
+    with pytest.raises(ValueError, match="same reference"):
+        load_prediction_yamls([y], str(reference))
+
+
+def test_load_prediction_yamls_rejects_mixed_coord(tmp_path: Path):
+    reference = tmp_path / "reference.parquet"
+    _reference_frame().collect().write_parquet(reference)
+    a = _write_prediction_yaml(tmp_path, reference, tag="a", coord="global")
+    b = _write_prediction_yaml(tmp_path, reference, tag="b", coord="local")
+    with pytest.raises(ValueError, match="coord"):
+        load_prediction_yamls([a, b], str(reference))
+
+
+def test_prep_with_prediction_writes_run_meta(tmp_path: Path):
+    rollout_yaml = _write_inputs(tmp_path)
+    reference = load_rollout_yaml(rollout_yaml)["dataset"]
+    pred_yaml = _write_prediction_yaml(tmp_path, Path(reference))
+    run_dir = _prep([rollout_yaml], prediction_yamls=[pred_yaml])
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    assert [p["name"] for p in meta.predictions] == ["prediction"]
+    assert meta.predictions[0]["plot_meta"]["checkpoint"] == "/ckpt/p.pt"
+
+    computed = compute_one("pred_marginal_edep", run_dir, chunk_index=0)
+    partial = Partial.load(computed)
+    assert partial.data["available"]
+
+
 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")

tests/test_render.py

@@ -329,6 +329,58 @@ def test_render_one_of_each_kind(tmp_path: Path):
                 "log_color": True,
             },
         ),
+        Reduced(
+            "ph1",
+            "prediction",
+            "paired_hist",
+            "Paired hist (single prediction)",
+            "x",
+            {"edges": [0, 1, 2, 3], "series": {"pred": {"pred": [1, 2, 3], "true": [2, 2, 2]}}, "log_y": False},
+        ),
+        Reduced(
+            "ph2",
+            "prediction",
+            "paired_hist",
+            "Paired hist (two predictions)",
+            "x",
+            {
+                "edges": [0, 1, 2, 3],
+                "series": {"a": {"pred": [1, 2, 3], "true": [2, 2, 2]}, "b": {"pred": [3, 2, 1]}},
+                "log_y": False,
+            },
+        ),
+        Reduced(
+            "hm2d",
+            "prediction",
+            "heatmap2d",
+            "Scatter (truth vs pred)",
+            "true x",
+            {
+                "x_edges": [0, 1, 2],
+                "y_edges": [0, 1, 2],
+                "series": {"pred": [[2, 0], [1, 3]]},
+                "ylabel": "predicted x",
+                "cbar_label": "count",
+                "log_color": True,
+                "diagonal": True,
+            },
+        ),
+        Reduced(
+            "profile_noref",
+            "prediction",
+            "profile",
+            "Residual profile (no reference)",
+            "true x",
+            {"edges": [0, 1, 2], "series": {"pred": {"mean": [0.1, -0.1], "std": [0.2, 0.2]}}},
+        ),
+        Reduced(
+            "bar_noref",
+            "prediction",
+            "bar",
+            "Constraint violations (no reference)",
+            "check",
+            {"labels": ["a", "b"], "series": {"pred": [0.01, 0.0]}, "ylabel": "rate"},
+        ),
     ]
     try:
         pdfs = _try_render(reduced, tmp_path)