diff --git a/giant/analysis/condor.py b/giant/analysis/condor.py index 039c6e9..60e1424 100644 --- a/giant/analysis/condor.py +++ b/giant/analysis/condor.py @@ -65,12 +65,22 @@ _PLOT_META_KEYS = ( "max_steps", "steps", "max_tracks_per_event", + "escape_threshold", + "n_events", "n_seed_events", "timestamp", "comment", + "weights", + "batch_size", + "device", + "rollout_seed", + "n_rows", + "termination_reason_counts", "model_config", "training_epoch", "best_val_loss", + "training_config", + "training_meta", ) diff --git a/giant/analysis/render.py b/giant/analysis/render.py index 76d2cd4..3941981 100644 --- a/giant/analysis/render.py +++ b/giant/analysis/render.py @@ -41,15 +41,46 @@ def _overlay(ax, edges: np.ndarray, series: dict[str, list], log_y: bool) -> Non ax.set_yscale("log") -def _nn_params(run_meta: dict) -> dict: - """Flatten the rollout's model/training provenance for the figure subtitle.""" - params = { - k: v for k, v in (run_meta.get("model_config") or {}).items() if v is not None - } +def _router_summary(model_config: dict) -> str: + r = model_config.get("router") or {} + if not r.get("enabled"): + return "off" + return f"{r.get('type', '?')}×{r.get('n_experts', '?')}" + + +def _figure_params(run_meta: dict) -> dict: + """Curated run identity for the figure subtitle (``new_figure(params=...)``). + + ``run_meta``/each plot's own ``.yaml`` (see ``_plot_metadata``) already + carry every threaded model/training/rollout/dataset parameter for + after-the-fact lookup — this picks only the handful that matter for + telling figures apart at a glance while flipping through a gallery, since + the subtitle is one unwrapped line of text. The last slot is + architecture-conditional: flow/ddpm runs show the ODE ``steps`` used for + this rollout, wgan runs show ``noise_dim`` instead since wgan sampling is + single-pass and has no ODE step count. + """ + mc = run_meta.get("model_config") or {} + mode = mc.get("mode") + params: dict = {} + if mc.get("hidden_dim") is not None: + params["hidden_dim"] = mc["hidden_dim"] + if mc.get("n_blocks") is not None: + params["n_blocks"] = mc["n_blocks"] + if mode is not None: + params["mode"] = mode + if mc.get("conditioning") is not None: + params["conditioning"] = mc["conditioning"] + params["router"] = _router_summary(mc) if run_meta.get("training_epoch") is not None: params["epoch"] = run_meta["training_epoch"] if run_meta.get("best_val_loss") is not None: params["best_val_loss"] = round(run_meta["best_val_loss"], 4) + if mode == "wgan": + if mc.get("noise_dim") is not None: + params["noise_dim"] = mc["noise_dim"] + elif run_meta.get("steps") is not None: + params["steps"] = run_meta["steps"] return params @@ -225,7 +256,7 @@ _RENDERERS = { def render(r: Reduced, run_meta: dict | None = None): """Build the matplotlib figure for one reduced artifact (dispatch on kind).""" - return _RENDERERS[r.kind](r, _nn_params(run_meta or {})) + return _RENDERERS[r.kind](r, _figure_params(run_meta or {})) def _plot_metadata(r: Reduced, run_meta: dict) -> dict: @@ -238,6 +269,11 @@ def _plot_metadata(r: Reduced, run_meta: dict) -> dict: meta.update(r.meta) if "note" in r.payload: meta["note"] = r.payload["note"] + if run_meta: + # Every threaded model/training/rollout/dataset parameter, so a + # single plot's metadata is self-contained for later comparison + # without cross-referencing the run's root metadata.yaml. + meta["parameters"] = {k: v for k, v in run_meta.items() if k != "title"} return meta diff --git a/giant/cli.py b/giant/cli.py index 511bad2..9ba234f 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -1024,6 +1024,9 @@ def rollout( ) raise typer.Exit(1) + gconfig.warn_if_checkpoint_config_mismatch(checkpoint) + training_cfg = gconfig.load_checkpoint_config(checkpoint) + model_cfg = ckpt["model_config"] conditioning = model_cfg.get("conditioning", "embedding") pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()} @@ -1104,17 +1107,26 @@ def rollout( "max_steps": max_steps, "steps": steps, "max_tracks_per_event": max_tracks_per_event, + "escape_threshold": escape_threshold, + "n_events": n_events, "n_seed_events": int(len(seeds["event_id"])), - "model_config": { - "mode": model_cfg.get("mode", "flow"), - "hidden_dim": model_cfg.get("hidden_dim"), - "n_blocks": model_cfg.get("n_blocks"), - "emb_dim": model_cfg.get("emb_dim"), - "dropout": model_cfg.get("dropout"), - "conditioning": conditioning, - }, + "weights": weights.value, + "batch_size": batch_size, + "device": str(_device), + "rollout_seed": seed, + "n_rows": summary["n_rows"], + "termination_reason_counts": summary["termination_reason_counts"], + # Full architecture spec baked into the checkpoint — includes the + # entire router sub-dict, not just a hand-picked subset, so any + # model knob (router type/n_experts, noise_dim, vocab sizes, ...) + # is available downstream without touching this command again. + "model_config": dict(model_cfg), "training_epoch": ckpt.get("epoch"), "best_val_loss": ckpt.get("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", {})), } ) ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False)) diff --git a/giant/config.py b/giant/config.py index cfd7a1b..2de0df5 100644 --- a/giant/config.py +++ b/giant/config.py @@ -188,6 +188,21 @@ def warn_if_git_hash_mismatch(file_cfg: dict, config_path: Path) -> None: ) +def load_checkpoint_config(ckpt_path: str | Path) -> dict: + """Load the full ``[train]``/``[model]``/``[meta]`` config.toml written + alongside a checkpoint by ``save_config``. + + Returns ``{}`` if no config.toml sits next to the checkpoint (older runs, + or a checkpoint moved without its sidecar) — this is best-effort + provenance for threading into a rollout's YAML sidecar, not a hard + requirement for using the checkpoint itself. + """ + config_path = Path(ckpt_path).parent / "config.toml" + if not config_path.exists(): + return {} + return load_toml(config_path) + + def warn_if_checkpoint_config_mismatch(ckpt_path: str | Path) -> None: """Look for a config.toml next to a checkpoint and warn on a git_hash mismatch.