analyze: thread full model/training/rollout/dataset params to plots

giant rollout now writes the checkpoint's complete model_config (incl.
the router sub-dict), the sibling config.toml's [train]/[meta] sections,
and every rollout CLI knob (weights, batch_size, escape_threshold,
n_events, device, seed) into the YAML sidecar instead of a hand-picked
subset. All of it flows through run_meta.json into each plot's own
metadata.yaml for later comparison, while the figure subtitle itself
shows a curated slice (hidden_dim, n_blocks, mode, conditioning, router,
epoch, best_val_loss, steps/noise_dim) via new_figure's params option.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

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

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 ``<id>.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
 
 

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

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.