diff --git a/CLAUDE.md b/CLAUDE.md index 912efaa..81fcdae 100644 --- a/CLAUDE.md +++ b/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 --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/.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_/`) 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/.json`; the local `giant analyze render ` 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. diff --git a/giant/analysis/__init__.py b/giant/analysis/__init__.py index 8a2f56a..a646f80 100644 --- a/giant/analysis/__init__.py +++ b/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", diff --git a/giant/analysis/condor.py b/giant/analysis/condor.py index cb34514..28cbcd8 100644 --- a/giant/analysis/condor.py +++ b/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: + + /shared.json fixed bin edges / group sets (prep) + /run_meta.json resolved rollout/reference paths + plot metadata + /reduced/.json one per compute job + /plots//.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/.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/.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 (``/analyze.sub``). Does - not submit — call ``condor_submit`` on the returned file, or use ``submit``. + Returns the submit description path (``/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 (``/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 diff --git a/giant/analysis/render.py b/giant/analysis/render.py index dee991a..a7d9014 100644 --- a/giant/analysis/render.py +++ b/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: ``/reduced`` → ``/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 + ) diff --git a/giant/analysis/sources.py b/giant/analysis/sources.py index e9dd6da..d71b129 100644 --- a/giant/analysis/sources.py +++ b/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) diff --git a/giant/cli.py b/giant/cli.py index 00b49ab..d989d08 100644 --- a/giant/cli.py +++ b/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") diff --git a/tests/test_condor.py b/tests/test_condor.py index c4c418b..ef2b89a 100644 --- a/tests/test_condor.py +++ b/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