analyze: drive prep/submit from the rollout YAML sidecar

`giant analyze prep` / `submit` now take the `giant rollout` YAML sidecar as
their only positional input instead of explicit --rollout/--reference/--out-dir.
The YAML's `output`/`dataset` keys name the rollout parquet and its seed file
(the reference truth), and the rest of the sidecar (checkpoint, geometry oracle,
cutoffs) flows into every plot's gallery metadata.

prep derives its own run directory next to the rollout parquet
(<...>/analysis_<id>/) holding shared.json, run_meta.json, reduced/, plots/.
compute-one and render now take just --run-dir / a run-dir argument and read the
resolved paths + metadata from run_meta.json, so the condor wrapper no longer
threads file paths. open_side scans a directory of reference shards via glob.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

CLAUDE.md

@@ -14,9 +14,8 @@ giant train path/to/steps.parquet --mode flow   # train (flow matching)
 giant train path/to/steps.parquet --mode ddpm   # train (DDPM baseline)
 giant predict path/to/steps.parquet --checkpoint ckpt/best.pt   # per-step predictions
 giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl  # full showers
-giant analyze submit --rollout roll.parquet --reference test.parquet --out-dir run/ \
-  --accounting-group cms                          # parallel rollout-vs-reference analysis on HTCondor
-giant analyze render --reduced-dir run/reduced --out run/plots --gallery  # render PDFs + HTML gallery
+giant analyze submit rollout.yaml --accounting-group cms   # parallel rollout-vs-reference analysis on HTCondor
+giant analyze render <run_dir> --gallery                   # render PDFs + HTML gallery (run_dir from prep/submit)
 dwarf --help                                     # dataset/tooling CLI: convert, migrate, bump-gen,
                                                   # bump-schema, status, update-manifest, create-manifest,
                                                   # make-root, build-geometry-oracle, hparam-scan
@@ -59,7 +58,7 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
 
 **Validation** (`giant/validate.py`): step-level marginal comparisons.
 
-**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one autoregressive `giant rollout` (for a given checkpoint) against a held-out miniCaloSim reference steps file, and produces publication-styled PDFs assembled into an HTML gallery. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json`, so every compute job is one pass with no range scan), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles, species/leakage, secondaries), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`). **Compute/render split:** `giant analyze submit` runs `prep` then submits one HTCondor job per plot (`compute-one`, polars/numpy only — no LaTeX on workers), each writing a small `reduced/<id>.json`; the local `giant analyze render` turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`.
+**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one autoregressive `giant rollout` (for a given checkpoint) against a held-out miniCaloSim reference steps file, and produces publication-styled PDFs assembled into an HTML gallery. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json`, so every compute job is one pass with no range scan), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles, species/leakage, secondaries), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`). **Input is a `giant rollout` YAML sidecar** (`condor.py:load_rollout_yaml`): its `output`/`dataset` keys name the rollout parquet and the seed file (= the reference truth), and the rest of the YAML (checkpoint, geometry oracle, cutoffs) flows into each plot's gallery metadata. `prep` derives its own **run directory** next to the rollout parquet (`<...>/analysis_<id>/`) holding `shared.json`, `run_meta.json`, `reduced/`, `plots/`. **Compute/render split:** `giant analyze submit rollout.yaml` runs `prep` then submits one HTCondor job per plot (`compute-one --run-dir`, polars/numpy only — no LaTeX on workers), each writing a small `reduced/<id>.json`; the local `giant analyze render <run_dir>` turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`.
 
 **Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — 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 energy cutoff, per-track max steps, escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
 

giant/analysis/__init__.py

@@ -10,7 +10,16 @@ re-exported here is plotstyle-free so it runs on a compute worker. Import
 """
 
 from giant.analysis.catalog import build_catalog, catalog_ids, get_spec
-from giant.analysis.condor import SubmitConfig, compute_one, prep, write_submit
+from giant.analysis.condor import (
+    RunMeta,
+    SubmitConfig,
+    compute_one,
+    compute_reduced,
+    derive_run_dir,
+    load_rollout_yaml,
+    prep,
+    write_submit,
+)
 from giant.analysis.context import Context, build_context
 from giant.analysis.reduced import Reduced
 from giant.analysis.sources import Side
@@ -19,8 +28,12 @@ __all__ = [
     "build_catalog",
     "catalog_ids",
     "get_spec",
+    "RunMeta",
     "SubmitConfig",
     "compute_one",
+    "compute_reduced",
+    "derive_run_dir",
+    "load_rollout_yaml",
     "prep",
     "write_submit",
     "Context",

giant/analysis/condor.py

@@ -1,15 +1,32 @@
-"""HTCondor orchestration: prep, per-plot compute, and the submit description.
+"""HTCondor orchestration driven by a ``giant rollout`` YAML sidecar.
+
+A rollout writes a YAML sidecar (``giant/cli.py:_write_prediction_ref`` +
+rollout extras) that already names both files we need and carries the run's
+provenance:
+
+* ``output``  — the rollout steps parquet (the *generated* side)
+* ``dataset`` — the file the rollout was seeded from, i.e. the held-out real
+  steps (the *reference* side)
+* ``checkpoint``, ``geometry_oracle``, ``energy_cutoff``, ``steps``, ... —
+  metadata that flows straight into every plot's gallery ``metadata.yaml``.
+
+So the analysis takes that one YAML as input, derives its own **run directory**
+next to the rollout parquet, and lays everything out under it:
+
+    <run_dir>/shared.json     fixed bin edges / group sets (prep)
+    <run_dir>/run_meta.json   resolved rollout/reference paths + plot metadata
+    <run_dir>/reduced/<id>.json   one per compute job
+    <run_dir>/plots/<family>/<id>.pdf   rendered locally
 
 Job model (one condor job per plot, compute/render split):
 
-1. ``prep`` runs once on the submit node — resolves the shared context (fixed bin
-   edges, energy quantiles, top species/materials) from a subsample and writes
-   ``shared.json``. Cheap; no LaTeX.
-2. one job per catalog id runs ``giant analyze compute-one`` on a worker — a
-   single streaming pass producing ``reduced/<id>.json``. polars/numpy only, no
-   LaTeX, so it needs no plotstyle in the container.
-3. a final *local* ``giant analyze render`` turns every reduced artifact into a
-   styled PDF + gallery metadata (that step imports plotstyle/LaTeX).
+1. ``prep`` runs once on the submit node — reads the YAML, resolves the shared
+   context from a subsample, writes ``shared.json`` + ``run_meta.json``.
+2. one job per catalog id runs ``giant analyze compute-one --run-dir`` on a
+   worker — a single streaming pass writing ``reduced/<id>.json`` (polars/numpy
+   only, no LaTeX).
+3. a final *local* ``giant analyze render`` turns those into the styled PDF +
+   gallery tree (that step imports plotstyle/LaTeX).
 
 Files on ``/ceph`` or ``/work`` are reached via ``ProvidesETPResources``; no
 HTCondor file transfer of the multi-GB inputs.
@@ -17,26 +34,123 @@ HTCondor file transfer of the multi-GB inputs.
 
 from __future__ import annotations
 
+import json
 from dataclasses import dataclass
 from pathlib import Path
 
+import yaml
+
 from giant.analysis.catalog import Bundle, catalog_ids, get_spec
 from giant.analysis.context import Context, build_context
 
+# Keys copied verbatim from a rollout YAML into each plot's gallery metadata.
+_PLOT_META_KEYS = (
+    "prediction_id",
+    "checkpoint",
+    "output",
+    "dataset",
+    "geometry_oracle",
+    "energy_cutoff",
+    "max_steps",
+    "steps",
+    "max_tracks_per_event",
+    "n_seed_events",
+    "timestamp",
+    "comment",
+)
+
+
+# ---------------------------------------------------------------------------
+# rollout-YAML → run directory
+# ---------------------------------------------------------------------------
+
+
+def load_rollout_yaml(path: str | Path) -> dict:
+    """Load a ``giant rollout`` YAML sidecar, requiring the two file paths."""
+    d = yaml.safe_load(Path(path).read_text())
+    for key in ("output", "dataset"):
+        if key not in d:
+            raise ValueError(
+                f"{path} is not a rollout YAML (missing {key!r}); expected the "
+                "sidecar `giant rollout` writes next to the checkpoint"
+            )
+    if d.get("kind") not in (None, "rollout"):
+        raise ValueError(f"{path} has kind={d.get('kind')!r}, not a rollout YAML")
+    return d
+
+
+def derive_run_dir(rollout_yaml: dict, run_dir: str | Path | None = None) -> Path:
+    """Analysis output directory, next to the rollout parquet unless overridden."""
+    if run_dir is not None:
+        return Path(run_dir)
+    rollout = Path(rollout_yaml["output"])
+    tag = str(rollout_yaml.get("prediction_id") or rollout.stem)[:8]
+    return rollout.parent / f"analysis_{tag}"
+
+
+def _plot_meta(rollout_yaml: dict) -> dict:
+    return {k: rollout_yaml[k] for k in _PLOT_META_KEYS if k in rollout_yaml}
+
+
+@dataclass
+class RunMeta:
+    """Resolved paths + plot metadata for one analysis run (``run_meta.json``)."""
+
+    rollout: str
+    reference: str
+    run_dir: str
+    title: str
+    plot_meta: dict
+
+    def save(self, path: str | Path) -> None:
+        Path(path).write_text(json.dumps(self.__dict__, indent=2))
+
+    @classmethod
+    def load(cls, path: str | Path) -> "RunMeta":
+        return cls(**json.loads(Path(path).read_text()))
+
+
+def prep(
+    rollout_yaml: str | Path,
+    run_dir: str | Path | None = None,
+    **ctx_kwargs,
+) -> Path:
+    """Read the rollout YAML, build the shared context, and lay out the run dir.
+
+    Writes ``shared.json`` + ``run_meta.json`` and returns the run directory.
+    """
+    y = load_rollout_yaml(rollout_yaml)
+    run_path = derive_run_dir(y, run_dir)
+    run_path.mkdir(parents=True, exist_ok=True)
+
+    rollout, reference = y["output"], y["dataset"]
+    ctx = build_context(rollout, reference, **ctx_kwargs)
+    ctx.save(run_path / "shared.json")
+
+    ckpt = Path(y.get("checkpoint", "")).name or "rollout"
+    RunMeta(
+        rollout=str(rollout),
+        reference=str(reference),
+        run_dir=str(run_path),
+        title=f"GIANT rollout analysis — {ckpt}",
+        plot_meta=_plot_meta(y),
+    ).save(run_path / "run_meta.json")
+    return run_path
+
 
 # ---------------------------------------------------------------------------
 # per-plot compute (what each condor job runs)
 # ---------------------------------------------------------------------------
 
 
-def compute_one(
+def compute_reduced(
     spec_id: str,
     rollout: str | Path,
     reference: str | Path,
     shared: str | Path,
     out: str | Path,
 ) -> Path:
-    """Run one plot's streaming reduction and write its ``Reduced`` JSON."""
+    """Core: run one plot's reduction against explicit paths → ``Reduced`` JSON."""
     ctx = Context.load(shared)
     bundle = Bundle.open(rollout, reference, ctx)
     reduced = get_spec(spec_id).compute(bundle)
@@ -45,6 +159,19 @@ def compute_one(
     return out
 
 
+def compute_one(spec_id: str, run_dir: str | Path) -> Path:
+    """Run one plot's reduction from a prepped run directory."""
+    run_path = Path(run_dir)
+    meta = RunMeta.load(run_path / "run_meta.json")
+    return compute_reduced(
+        spec_id,
+        meta.rollout,
+        meta.reference,
+        run_path / "shared.json",
+        run_path / "reduced" / f"{spec_id}.json",
+    )
+
+
 # ---------------------------------------------------------------------------
 # submit description
 # ---------------------------------------------------------------------------
@@ -52,9 +179,7 @@ def compute_one(
 
 @dataclass
 class SubmitConfig:
-    rollout: Path
-    reference: Path
-    out_dir: Path
+    run_dir: Path
     accounting_group: str
     repo_dir: Path
     docker_image: str = "mschnepf/slc7-condocker"
@@ -67,18 +192,13 @@ class SubmitConfig:
 _WRAPPER = """#!/bin/bash
 set -euo pipefail
 cd {repo_dir}
-exec uv run giant analyze compute-one \\
-  --id "$1" \\
-  --rollout {rollout} \\
-  --reference {reference} \\
-  --shared {shared} \\
-  --out {reduced_dir}/"$1".json
+exec uv run giant analyze compute-one --id "$1" --run-dir {run_dir}
 """
 
 
 def _submit_description(cfg: SubmitConfig, wrapper: Path, ids_file: Path) -> str:
     reqs_attrs = (
-        "+RemoteJob = True\nrequest_walltime = {wt}\n".format(wt=cfg.request_walltime_s)
+        "+RemoteJob = True\n"
         if cfg.remote
         else "requirements = TARGET.ProvidesETPResources\n"
     )
@@ -94,9 +214,9 @@ def _submit_description(cfg: SubmitConfig, wrapper: Path, ids_file: Path) -> str
         f"+RequestWalltime = {cfg.request_walltime_s}\n"
         f"accounting_group = {cfg.accounting_group}\n"
         f"{reqs_attrs}"
-        "output = logs/$(plotid).out\n"
-        "error = logs/$(plotid).err\n"
-        "log = logs/condor.log\n"
+        f"output = {cfg.run_dir}/logs/$(plotid).out\n"
+        f"error = {cfg.run_dir}/logs/$(plotid).err\n"
+        f"log = {cfg.run_dir}/logs/condor.log\n"
         f"queue plotid from {ids_file}\n"
     )
 
@@ -104,45 +224,21 @@ def _submit_description(cfg: SubmitConfig, wrapper: Path, ids_file: Path) -> str
 def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
     """Write the wrapper script, plot-id list, and HTCondor submit description.
 
-    Returns the path to the submit description (``<out_dir>/analyze.sub``). Does
-    not submit — call ``condor_submit`` on the returned file, or use ``submit``.
+    Returns the submit description path (``<run_dir>/analyze.sub``). Does not
+    submit — call ``condor_submit`` on the returned file.
     """
     ids = ids or catalog_ids()
-    out_dir = cfg.out_dir
-    reduced_dir = out_dir / "reduced"
-    (out_dir / "logs").mkdir(parents=True, exist_ok=True)
-    reduced_dir.mkdir(parents=True, exist_ok=True)
+    run_dir = cfg.run_dir
+    (run_dir / "logs").mkdir(parents=True, exist_ok=True)
+    (run_dir / "reduced").mkdir(parents=True, exist_ok=True)
 
-    wrapper = out_dir / "run_compute.sh"
-    wrapper.write_text(
-        _WRAPPER.format(
-            repo_dir=cfg.repo_dir,
-            rollout=cfg.rollout,
-            reference=cfg.reference,
-            shared=out_dir / "shared.json",
-            reduced_dir=reduced_dir,
-        )
-    )
+    wrapper = run_dir / "run_compute.sh"
+    wrapper.write_text(_WRAPPER.format(repo_dir=cfg.repo_dir, run_dir=run_dir))
     wrapper.chmod(0o755)
 
-    ids_file = out_dir / "plotids.txt"
+    ids_file = run_dir / "plotids.txt"
     ids_file.write_text("\n".join(ids) + "\n")
 
-    sub = out_dir / "analyze.sub"
+    sub = run_dir / "analyze.sub"
     sub.write_text(_submit_description(cfg, wrapper, ids_file))
     return sub
-
-
-def prep(
-    rollout: str | Path,
-    reference: str | Path,
-    out_dir: str | Path,
-    **kwargs,
-) -> Path:
-    """Build and save the shared context (``<out_dir>/shared.json``)."""
-    out_dir = Path(out_dir)
-    out_dir.mkdir(parents=True, exist_ok=True)
-    ctx = build_context(rollout, reference, **kwargs)
-    shared = out_dir / "shared.json"
-    ctx.save(shared)
-    return shared

giant/analysis/render.py

@@ -189,7 +189,7 @@ def render_all(
                 "title": run_meta.get("title", "GIANT rollout analysis"),
                 "description": "Autoregressive rollout compared against held-out Geant4 reference steps.",
                 "experiment": "GIANT",
-                "parameters": run_meta,
+                "parameters": {k: v for k, v in run_meta.items() if k != "title"},
             },
             sort_keys=False,
         )
@@ -204,3 +204,24 @@ def render_all(
     if run_gallery:
         subprocess.run(["gallery", "generate", "--source", str(out_dir)], check=True)
     return pdfs
+
+
+def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
+    """Render a prepped run directory: ``<run_dir>/reduced`` → ``<run_dir>/plots``.
+
+    Pulls the rollout provenance (checkpoint, paths, cutoffs) from
+    ``run_meta.json`` into every plot's gallery metadata.
+    """
+    from giant.analysis.condor import RunMeta
+
+    run_dir = Path(run_dir)
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    run_meta = {
+        "title": meta.title,
+        "rollout": meta.rollout,
+        "reference": meta.reference,
+        **meta.plot_meta,
+    }
+    return render_all(
+        run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery
+    )

giant/analysis/sources.py

@@ -112,6 +112,11 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
     path = Path(source)
     if side is Side.rollout:
         _check_rollout_metadata(path)
+        return pl.scan_parquet(path)
+    # The reference (a rollout's seed `dataset`) may be a directory of parquet
+    # shards rather than a single file — scan them all.
+    if path.is_dir():
+        return pl.scan_parquet(str(path / "**/*.parquet"))
     return pl.scan_parquet(path)
 
 

giant/cli.py

@@ -1123,32 +1123,35 @@ app.add_typer(analyze_app, name="analyze")
 
 @analyze_app.command("prep")
 def analyze_prep(
-    rollout: Annotated[Path, typer.Option("--rollout", help="giant rollout parquet")],
-    reference: Annotated[
-        Path, typer.Option("--reference", help="Reference miniCaloSim steps parquet")
-    ],
-    out_dir: Annotated[
+    rollout_yaml: Annotated[
         Path,
-        typer.Option(
-            "--out-dir", "-o", help="Run directory for shared.json / reduced / plots"
+        typer.Argument(
+            help="giant rollout YAML sidecar (names the rollout + reference files)"
         ),
     ],
+    run_dir: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--run-dir",
+            "-o",
+            help="Override the run directory (default: next to the rollout parquet)",
+        ),
+    ] = None,
     n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4,
     n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50,
     top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6,
 ) -> None:
-    """Resolve the shared context (fixed bin edges / group sets) → shared.json."""
+    """Read the rollout YAML → shared.json + run_meta.json in the run directory."""
     from giant.analysis import prep
 
-    shared = prep(
-        rollout,
-        reference,
-        out_dir,
+    path = prep(
+        rollout_yaml,
+        run_dir,
         n_energy_bins=n_energy_bins,
         n_marginal_bins=n_marginal_bins,
         top_k_pdg=top_k_pdg,
     )
-    typer.echo(f"wrote {shared}")
+    typer.echo(f"run directory: {path}")
 
 
 @analyze_app.command("compute-one")
@@ -1156,17 +1159,14 @@ def analyze_compute_one(
     id: Annotated[
         str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")
     ],
-    rollout: Annotated[Path, typer.Option("--rollout")],
-    reference: Annotated[Path, typer.Option("--reference")],
-    shared: Annotated[
-        Path, typer.Option("--shared", help="shared.json from `analyze prep`")
+    run_dir: Annotated[
+        Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")
     ],
-    out: Annotated[Path, typer.Option("--out", help="Output reduced JSON path")],
 ) -> None:
     """Run one plot's streaming reduction (this is what each condor job runs)."""
     from giant.analysis import compute_one
 
-    path = compute_one(id, rollout, reference, shared, out)
+    path = compute_one(id, run_dir)
     typer.echo(f"wrote {path}")
 
 
@@ -1181,10 +1181,9 @@ def analyze_list() -> None:
 
 @analyze_app.command("render")
 def analyze_render(
-    reduced_dir: Annotated[
-        Path, typer.Option("--reduced-dir", help="Directory of reduced *.json")
+    run_dir: Annotated[
+        Path, typer.Argument(help="Run directory from `analyze prep` (holds reduced/)")
     ],
-    out: Annotated[Path, typer.Option("--out", "-o", help="Output PDF/gallery tree")],
     gallery: Annotated[
         bool,
         typer.Option(
@@ -1193,18 +1192,20 @@ def analyze_render(
     ] = False,
 ) -> None:
     """Render reduced artifacts to styled PDFs + gallery metadata (local; needs LaTeX)."""
-    from giant.analysis.render import render_all
+    from giant.analysis.render import render_run
 
-    pdfs = render_all(reduced_dir, out, run_gallery=gallery)
-    typer.echo(f"rendered {len(pdfs)} plots → {out}")
+    pdfs = render_run(run_dir, run_gallery=gallery)
+    typer.echo(f"rendered {len(pdfs)} plots → {Path(run_dir) / 'plots'}")
 
 
 @analyze_app.command("submit")
 def analyze_submit(
-    rollout: Annotated[Path, typer.Option("--rollout")],
-    reference: Annotated[Path, typer.Option("--reference")],
-    out_dir: Annotated[Path, typer.Option("--out-dir", "-o")],
+    rollout_yaml: Annotated[Path, typer.Argument(help="giant rollout YAML sidecar")],
     accounting_group: Annotated[str, typer.Option("--accounting-group")],
+    run_dir: Annotated[
+        Optional[Path],
+        typer.Option("--run-dir", "-o", help="Override the run directory"),
+    ] = None,
     docker_image: Annotated[
         str, typer.Option("--docker-image")
     ] = "mschnepf/slc7-condocker",
@@ -1222,11 +1223,9 @@ def analyze_submit(
 
     from giant.analysis import SubmitConfig, prep, write_submit
 
-    prep(rollout, reference, out_dir)
+    path = prep(rollout_yaml, run_dir)
     cfg = SubmitConfig(
-        rollout=rollout.resolve(),
-        reference=reference.resolve(),
-        out_dir=out_dir,
+        run_dir=path,
         accounting_group=accounting_group,
         repo_dir=Path.cwd(),
         docker_image=docker_image,
@@ -1234,6 +1233,7 @@ def analyze_submit(
         remote=remote,
     )
     sub = write_submit(cfg)
+    typer.echo(f"run directory: {path}")
     typer.echo(f"wrote submit description: {sub}")
     if dry_run:
         typer.echo("dry-run: not submitting")

tests/test_condor.py

@@ -1,85 +1,127 @@
-"""Tests for the HTCondor submit description + the compute_one round-trip."""
+"""Tests for the rollout-YAML → run-directory flow, compute, and submit."""
 
 from __future__ import annotations
 
 from pathlib import Path
 
+import pyarrow.parquet as pq
+import pytest
+import yaml
+
 from giant.analysis import (
-    Context,
+    RunMeta,
     SubmitConfig,
-    build_context,
     catalog_ids,
     compute_one,
+    compute_reduced,
+    derive_run_dir,
+    load_rollout_yaml,
     prep,
     write_submit,
 )
+from giant.analysis.condor import Context
 from giant.analysis.reduced import Reduced
+from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE
 from tests.test_analysis_reduce import _reference_frame, _rollout_frame
 
 
-def test_prep_and_compute_one_roundtrip(tmp_path: Path):
-    r, t = _rollout_frame(), _reference_frame()
-    ctx = build_context(
-        r, t, n_energy_bins=2, n_marginal_bins=8, top_k_pdg=3, sample_rows=1000
-    )
-    shared = tmp_path / "shared.json"
-    ctx.save(shared)
-    assert Context.load(shared).top_pdgs == ctx.top_pdgs
+def _write_inputs(tmp_path: Path) -> Path:
+    """Materialize rollout+reference parquet and a rollout YAML; return the YAML path."""
+    rollout = tmp_path / "rollout.parquet"
+    reference = tmp_path / "reference.parquet"
+    tbl = _rollout_frame().collect().to_arrow()
+    tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE})
+    pq.write_table(tbl, rollout)
+    _reference_frame().collect().write_parquet(reference)
 
-    out = compute_one("marginal_edep", r, t, shared, tmp_path / "marginal_edep.json")
+    yaml_path = tmp_path / "run.yaml"
+    yaml_path.write_text(
+        yaml.safe_dump(
+            {
+                "prediction_id": "abcd1234ef",
+                "output": str(rollout),
+                "dataset": str(reference),
+                "checkpoint": "/ckpt/best.pt",
+                "kind": "rollout",
+                "energy_cutoff": 0.1,
+                "steps": 10,
+            }
+        )
+    )
+    return yaml_path
+
+
+_CTX = dict(n_energy_bins=2, n_marginal_bins=8, top_k_pdg=3, sample_rows=1000)
+
+
+def test_load_rollout_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_rollout_yaml(bad)
+
+
+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")
+    assert derive_run_dir(y, "/somewhere") == Path("/somewhere")
+
+
+def test_prep_lays_out_run_dir(tmp_path: Path):
+    yaml_path = _write_inputs(tmp_path)
+    run_dir = prep(yaml_path, **_CTX)
+    assert run_dir == tmp_path / "analysis_abcd1234"
+    assert (run_dir / "shared.json").exists()
+    ctx = Context.load(run_dir / "shared.json")
+    assert set(ctx.var_ranges) == {"step_length", "edep", "delta_e", "post_E"}
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    assert meta.reference.endswith("reference.parquet")
+    assert meta.plot_meta["checkpoint"] == "/ckpt/best.pt"
+    assert "best.pt" in meta.title
+
+
+def test_compute_one_from_run_dir(tmp_path: Path):
+    run_dir = prep(_write_inputs(tmp_path), **_CTX)
+    out = compute_one("marginal_edep", run_dir)
+    assert out == run_dir / "reduced" / "marginal_edep.json"
     reduced = Reduced.load(out)
     assert reduced.id == "marginal_edep"
     assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1
 
 
-def test_prep_writes_shared_json(tmp_path: Path):
-    r, t = _rollout_frame(), _reference_frame()
-    shared = prep(
-        r,
-        t,
-        tmp_path / "run",
-        n_energy_bins=2,
-        n_marginal_bins=8,
-        top_k_pdg=3,
-        sample_rows=1000,
+def test_compute_reduced_explicit_paths(tmp_path: Path):
+    run_dir = prep(_write_inputs(tmp_path), **_CTX)
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    out = compute_reduced(
+        "marginal_step_length",
+        meta.rollout,
+        meta.reference,
+        run_dir / "shared.json",
+        tmp_path / "r.json",
     )
-    assert shared.exists()
-    ctx = Context.load(shared)
-    assert set(ctx.var_ranges) == {"step_length", "edep", "delta_e", "post_E"}
+    assert Reduced.load(out).id == "marginal_step_length"
 
 
 def test_write_submit_description(tmp_path: Path):
-    cfg = SubmitConfig(
-        rollout=tmp_path / "r.parquet",
-        reference=tmp_path / "t.parquet",
-        out_dir=tmp_path / "run",
-        accounting_group="cms",
-        repo_dir=tmp_path,
-    )
-    sub = write_submit(cfg)
-    txt = sub.read_text()
+    run_dir = prep(_write_inputs(tmp_path), **_CTX)
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path)
+    txt = write_submit(cfg).read_text()
     assert "universe = docker" in txt
     assert "docker_image = mschnepf/slc7-condocker" in txt
     assert "requirements = TARGET.ProvidesETPResources" in txt
     assert "accounting_group = cms" in txt
     assert "queue plotid from" in txt
-    # one queue item per catalog id
-    ids = (cfg.out_dir / "plotids.txt").read_text().split()
-    assert ids == catalog_ids()
-    # wrapper is executable and self-contained
-    wrapper = cfg.out_dir / "run_compute.sh"
+    assert (run_dir / "plotids.txt").read_text().split() == catalog_ids()
+    wrapper = run_dir / "run_compute.sh"
     assert wrapper.exists() and (wrapper.stat().st_mode & 0o111)
-    assert "giant analyze compute-one" in wrapper.read_text()
+    body = wrapper.read_text()
+    assert "giant analyze compute-one --id" in body and "--run-dir" in body
 
 
 def test_write_submit_remote_flag(tmp_path: Path):
+    run_dir = prep(_write_inputs(tmp_path), **_CTX)
     cfg = SubmitConfig(
-        rollout=tmp_path / "r.parquet",
-        reference=tmp_path / "t.parquet",
-        out_dir=tmp_path / "run",
-        accounting_group="cms",
-        repo_dir=tmp_path,
-        remote=True,
+        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, remote=True
     )
     txt = write_submit(cfg).read_text()
     assert "+RemoteJob = True" in txt