#!/usr/bin/env python3 from __future__ import annotations import argparse import copy import csv import json import os import sys from dataclasses import asdict, dataclass, replace from pathlib import Path REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") MPLCONFIGDIR = REPO_ROOT / ".cache" / "matplotlib" MPLCONFIGDIR.mkdir(parents=True, exist_ok=True) os.environ.setdefault("MPLCONFIGDIR", str(MPLCONFIGDIR)) os.environ.setdefault("MUJOCO_GL", "egl") import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from matplotlib.collections import LineCollection from matplotlib.colors import Normalize, TwoSlopeNorm import numpy as np import torch from common import ensure_repo_on_path ensure_repo_on_path() from hamworld.runtime import make_env from hamworld.world_model import CanonicalDynamicsWorldModel, infer_checkpoint_step PAPER_TITLE_FONTSIZE = 19 PAPER_LABEL_FONTSIZE = 18 PAPER_TICK_FONTSIZE = 15 PAPER_CBAR_FONTSIZE = 17 PAPER_LEGEND_FONTSIZE = 14 PAPER_SPINE_WIDTH = 1.2 @dataclass(frozen=True) class PairScore: q_index: int p_index: int kind: str score: float coverage: float swirl: float speed: float occupancy: float q_span: float p_span: float @dataclass(frozen=True) class PairSummary: q_index: int p_index: int kind: str score: float num_segments: int mean_abs_delta_h: float mean_control_norm: float mean_abs_control_energy_push: float mean_cross_ratio: float mean_abs_slice_cross_energy: float cross_threshold: float threshold_cross_rate: float corr_abs_delta_h_control_norm: float corr_abs_delta_h_abs_control_energy_push: float corr_cross_ratio_control_norm: float corr_cross_ratio_abs_control_energy_push: float corr_signed_delta_h_control_energy_push: float corr_abs_slice_cross_energy_abs_control_energy_push: float corr_binary_cross_abs_control_energy_push: float mean_cross_ratio_high_control: float mean_cross_ratio_low_control: float mean_abs_delta_h_high_control: float mean_abs_delta_h_low_control: float mean_cross_ratio_high_push: float mean_cross_ratio_low_push: float mean_abs_delta_h_high_push: float mean_abs_delta_h_low_push: float threshold_cross_rate_high_push: float threshold_cross_rate_low_push: float sweep_lift_auc: float sweep_binary_corr_auc: float sweep_best_threshold_quantile: float sweep_best_lift: float sweep_best_binary_corr: float def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Render Hamiltonian slice plots with control/crossing overlays.") parser.add_argument("--checkpoint", required=True, help="Path to checkpoint_*.pt.") parser.add_argument("--trace-path", required=True, help="Path to dynamics trace .npz file.") parser.add_argument("--output-dir", required=True, help="Directory for figures and summaries.") parser.add_argument("--device", default="cpu", help="Torch device, e.g. cpu or cuda:0.") parser.add_argument("--grid-size", type=int, default=81, help="Grid size for slice evaluation.") parser.add_argument("--same-k", type=int, default=4, help="How many same-index pairs to render.") parser.add_argument("--cross-k", type=int, default=2, help="How many cross-index pairs to render.") parser.add_argument("--trace-quantile-lo", type=float, default=2.0, help="Low percentile for slice bounds.") parser.add_argument("--trace-quantile-hi", type=float, default=98.0, help="High percentile for slice bounds.") parser.add_argument("--scatter-max-points", type=int, default=1200, help="Max points per scatter panel.") parser.add_argument( "--cross-threshold-quantile", type=float, default=0.60, help="Quantile of |∇H_slice·Δx| used to treat tiny crossings as no crossing.", ) parser.add_argument( "--cross-threshold-sweep", default="0.40,0.50,0.60,0.70,0.80,0.90", help="Comma-separated quantiles used for threshold sweep plots and summary metrics.", ) parser.add_argument("--focus-q-index", type=int, default=2, help="Focused q index for the summary sheet.") parser.add_argument("--focus-p-index", type=int, default=2, help="Focused p index for the summary sheet.") parser.add_argument("--ranking-top-k", type=int, default=12, help="How many top pairs to show in ranking charts.") return parser.parse_args() def _load_checkpoint_world_model(checkpoint_path: Path, device: torch.device) -> CanonicalDynamicsWorldModel: payload = torch.load(checkpoint_path, map_location="cpu", weights_only=False) config = copy.deepcopy(payload["config"]) env, spec = make_env(config["task"], int(config["experiment"]["seed"]) + 1000) close_fn = getattr(env, "close", None) if callable(close_fn): close_fn() world_model = CanonicalDynamicsWorldModel( config=config, obs_dim=int(spec.observation_shape[0]), action_dim=int(spec.action_shape[0]), ).to(device) agent_state = payload.get("agent_state") if isinstance(agent_state, dict) and "model" in agent_state: state_dict = agent_state["model"] else: state_dict = payload.get("model") if state_dict is None: raise ValueError(f"Checkpoint does not contain a world model state: {checkpoint_path}") world_model.load_state_dict(state_dict) world_model.set_schedule_step(infer_checkpoint_step(payload, checkpoint_path)) world_model.eval() return world_model def _load_trace(trace_path: Path) -> dict[str, np.ndarray]: trace = np.load(trace_path, allow_pickle=True) return {key: trace[key] for key in trace.files} def _flatten_valid_vectors(values: np.ndarray, valid_mask: np.ndarray) -> np.ndarray: mask = np.asarray(valid_mask, dtype=bool) return np.asarray(values[mask], dtype=np.float32) def _robust_span(values: np.ndarray, lo: float = 5.0, hi: float = 95.0) -> float: low, high = np.percentile(values, [lo, hi]) return float(high - low) def _histogram_entropy(x: np.ndarray, y: np.ndarray, bins: int = 20) -> float: hist, _, _ = np.histogram2d(x, y, bins=bins) probs = hist.ravel() total = probs.sum() if total <= 0.0: return 0.0 probs = probs[probs > 0.0] / total entropy = -(probs * np.log(probs)).sum() max_entropy = np.log(float(bins * bins)) return float(entropy / max(1e-6, max_entropy)) def _normalize(values: list[float]) -> list[float]: maximum = max(values) if values else 1.0 if maximum <= 1e-6: return [0.0 for _ in values] return [float(value / maximum) for value in values] def _parse_float_list(raw: str) -> list[float]: values: list[float] = [] for part in raw.split(","): item = part.strip() if not item: continue values.append(float(item)) if not values: raise ValueError("Expected at least one threshold quantile.") values = sorted(set(values)) for value in values: if not (0.0 < value < 1.0): raise ValueError(f"Threshold quantiles must be in (0, 1): {value}") return values def score_pairs(trace: dict[str, np.ndarray]) -> list[PairScore]: valid_mask = np.asarray(trace["valid_mask"], dtype=bool) q = _flatten_valid_vectors(trace["q"], valid_mask) p = _flatten_valid_vectors(trace["p"], valid_mask) dH_dq = _flatten_valid_vectors(trace["dH_dq"], valid_mask) dH_dp = _flatten_valid_vectors(trace["dH_dp"], valid_mask) q_dim = int(np.asarray(trace["q_dim"]).item()) p_dim = int(np.asarray(trace["p_dim"]).item()) rows: list[dict[str, float | int | str]] = [] for q_index in range(q_dim): for p_index in range(p_dim): qv = q[:, q_index] pv = p[:, p_index] dq = dH_dp[:, q_index] dp = -dH_dq[:, p_index] q_center = float(np.median(qv)) p_center = float(np.median(pv)) radial_norm = np.hypot(qv - q_center, pv - p_center) flow_norm = np.hypot(dq, dp) tangential = np.abs((qv - q_center) * dp - (pv - p_center) * dq) swirl = float(np.mean(tangential / np.maximum(1e-6, radial_norm * flow_norm))) q_span = _robust_span(qv) p_span = _robust_span(pv) rows.append( { "q_index": q_index, "p_index": p_index, "kind": "same" if q_index == p_index else "cross", "coverage": float(q_span * p_span), "swirl": swirl, "speed": float(np.mean(flow_norm)), "occupancy": _histogram_entropy(qv, pv), "q_span": q_span, "p_span": p_span, } ) coverage_norm = _normalize([float(row["coverage"]) for row in rows]) swirl_norm = _normalize([float(row["swirl"]) for row in rows]) speed_norm = _normalize([float(row["speed"]) for row in rows]) occupancy_norm = _normalize([float(row["occupancy"]) for row in rows]) scored: list[PairScore] = [] for index, row in enumerate(rows): score = ( 0.35 * coverage_norm[index] + 0.30 * swirl_norm[index] + 0.20 * speed_norm[index] + 0.15 * occupancy_norm[index] ) scored.append( PairScore( q_index=int(row["q_index"]), p_index=int(row["p_index"]), kind=str(row["kind"]), score=float(score), coverage=float(row["coverage"]), swirl=float(row["swirl"]), speed=float(row["speed"]), occupancy=float(row["occupancy"]), q_span=float(row["q_span"]), p_span=float(row["p_span"]), ) ) return sorted(scored, key=lambda item: item.score, reverse=True) def select_pairs(scores: list[PairScore], same_k: int, cross_k: int) -> list[PairScore]: same_pairs = [row for row in scores if row.kind == "same"][: max(0, int(same_k))] cross_pairs = [row for row in scores if row.kind == "cross"][: max(0, int(cross_k))] return same_pairs + cross_pairs def _slice_bounds(values: np.ndarray, quantile_lo: float, quantile_hi: float) -> tuple[float, float]: low, high = np.percentile(values, [quantile_lo, quantile_hi]) if not np.isfinite(low) or not np.isfinite(high) or abs(high - low) < 1e-5: center = float(np.median(values)) return center - 1.0, center + 1.0 margin = 0.18 * float(high - low) return float(low - margin), float(high + margin) def evaluate_slice( world_model: CanonicalDynamicsWorldModel, q_reference: np.ndarray, p_reference: np.ndarray, pair: PairScore, q_values: np.ndarray, p_values: np.ndarray, grid_size: int, quantile_lo: float, quantile_hi: float, device: torch.device, ) -> dict[str, np.ndarray]: x_min, x_max = _slice_bounds(q_values, quantile_lo, quantile_hi) y_min, y_max = _slice_bounds(p_values, quantile_lo, quantile_hi) q_axis = np.linspace(x_min, x_max, grid_size, dtype=np.float64) p_axis = np.linspace(y_min, y_max, grid_size, dtype=np.float64) q_grid, p_grid = np.meshgrid(q_axis, p_axis, indexing="xy") q_eval = np.repeat(q_reference[None, :], grid_size * grid_size, axis=0) p_eval = np.repeat(p_reference[None, :], grid_size * grid_size, axis=0) q_eval[:, pair.q_index] = q_grid.reshape(-1) p_eval[:, pair.p_index] = p_grid.reshape(-1) q_tensor = torch.as_tensor(q_eval, dtype=torch.float32, device=device).requires_grad_(True) p_tensor = torch.as_tensor(p_eval, dtype=torch.float32, device=device).requires_grad_(True) energy = world_model.energy_head(q_tensor, p_tensor) dH_dq, dH_dp = torch.autograd.grad(energy.sum(), (q_tensor, p_tensor), create_graph=False, retain_graph=False) h = energy.detach().cpu().numpy().reshape(grid_size, grid_size) u = dH_dp.detach().cpu().numpy()[:, pair.q_index].reshape(grid_size, grid_size) v = (-dH_dq.detach().cpu().numpy()[:, pair.p_index]).reshape(grid_size, grid_size) return { "q_axis": q_axis, "p_axis": p_axis, "H": h, "U": u, "V": v, "speed": np.hypot(u, v), } def _safe_pearson(x: np.ndarray, y: np.ndarray) -> float: x = np.asarray(x, dtype=np.float64) y = np.asarray(y, dtype=np.float64) if x.size < 2 or y.size < 2: return float("nan") x_std = np.std(x) y_std = np.std(y) if x_std <= 1e-12 or y_std <= 1e-12: return float("nan") return float(np.corrcoef(x, y)[0, 1]) def _quantile_mean(values: np.ndarray, selector: np.ndarray) -> float: if selector.size == 0 or not np.any(selector): return float("nan") return float(np.mean(values[selector])) def _quantile_mask(values: np.ndarray, quantile: float, upper: bool) -> np.ndarray: if values.size == 0: return np.zeros_like(values, dtype=bool) threshold = float(np.quantile(values, quantile)) if upper: return values >= threshold return values <= threshold def build_pair_steps(trace: dict[str, np.ndarray], pair: PairScore) -> dict[str, np.ndarray]: valid = np.asarray(trace["valid_mask"], dtype=bool) q = np.asarray(trace["q"], dtype=np.float32) p = np.asarray(trace["p"], dtype=np.float32) h = np.asarray(trace["H"], dtype=np.float32) dH_dq = np.asarray(trace["dH_dq"], dtype=np.float32) dH_dp = np.asarray(trace["dH_dp"], dtype=np.float32) control = np.asarray(trace["control"], dtype=np.float32) segments: list[np.ndarray] = [] x0s: list[float] = [] y0s: list[float] = [] x1s: list[float] = [] y1s: list[float] = [] delta_hs: list[float] = [] abs_delta_hs: list[float] = [] control_norms: list[float] = [] control_pushes: list[float] = [] abs_control_pushes: list[float] = [] cross_ratios: list[float] = [] abs_slice_cross_energies: list[float] = [] signed_slice_cross_energies: list[float] = [] step_norms: list[float] = [] episode_ids: list[int] = [] step_ids: list[int] = [] eps = 1e-8 for episode_index in range(q.shape[0]): length = q.shape[1] for step_index in range(length - 1): if not (valid[episode_index, step_index] and valid[episode_index, step_index + 1]): continue x0 = float(q[episode_index, step_index, pair.q_index]) y0 = float(p[episode_index, step_index, pair.p_index]) x1 = float(q[episode_index, step_index + 1, pair.q_index]) y1 = float(p[episode_index, step_index + 1, pair.p_index]) delta_h = float(h[episode_index, step_index + 1] - h[episode_index, step_index]) control_vec = control[episode_index, step_index] grad_q_i = float(dH_dq[episode_index, step_index, pair.q_index]) grad_p_i = float(dH_dp[episode_index, step_index, pair.p_index]) move = np.asarray([x1 - x0, y1 - y0], dtype=np.float64) normal = np.asarray([grad_q_i, grad_p_i], dtype=np.float64) cross_ratio = abs(float(move.dot(normal))) / max(eps, float(np.linalg.norm(move) * np.linalg.norm(normal))) control_push = float(np.dot(dH_dp[episode_index, step_index], control_vec)) signed_slice_cross_energy = float(move.dot(normal)) segments.append(np.asarray([[x0, y0], [x1, y1]], dtype=np.float64)) x0s.append(x0) y0s.append(y0) x1s.append(x1) y1s.append(y1) delta_hs.append(delta_h) abs_delta_hs.append(abs(delta_h)) control_norms.append(float(np.linalg.norm(control_vec))) control_pushes.append(control_push) abs_control_pushes.append(abs(control_push)) cross_ratios.append(min(1.0, max(0.0, cross_ratio))) abs_slice_cross_energies.append(abs(signed_slice_cross_energy)) signed_slice_cross_energies.append(signed_slice_cross_energy) step_norms.append(float(np.linalg.norm(move))) episode_ids.append(episode_index) step_ids.append(step_index) return { "segments": np.asarray(segments, dtype=np.float64), "x0": np.asarray(x0s, dtype=np.float64), "y0": np.asarray(y0s, dtype=np.float64), "x1": np.asarray(x1s, dtype=np.float64), "y1": np.asarray(y1s, dtype=np.float64), "delta_h": np.asarray(delta_hs, dtype=np.float64), "abs_delta_h": np.asarray(abs_delta_hs, dtype=np.float64), "control_norm": np.asarray(control_norms, dtype=np.float64), "control_energy_push": np.asarray(control_pushes, dtype=np.float64), "abs_control_energy_push": np.asarray(abs_control_pushes, dtype=np.float64), "cross_ratio": np.asarray(cross_ratios, dtype=np.float64), "abs_slice_cross_energy": np.asarray(abs_slice_cross_energies, dtype=np.float64), "signed_slice_cross_energy": np.asarray(signed_slice_cross_energies, dtype=np.float64), "step_norm": np.asarray(step_norms, dtype=np.float64), "episode_id": np.asarray(episode_ids, dtype=np.int32), "step_id": np.asarray(step_ids, dtype=np.int32), } def summarize_pair(pair: PairScore, steps: dict[str, np.ndarray], cross_threshold_quantile: float) -> PairSummary: control_norm = steps["control_norm"] abs_push = steps["abs_control_energy_push"] abs_delta_h = steps["abs_delta_h"] cross_ratio = steps["cross_ratio"] signed_delta_h = steps["delta_h"] push = steps["control_energy_push"] abs_slice_cross_energy = steps["abs_slice_cross_energy"] control_hi = np.quantile(control_norm, 0.75) if control_norm.size else float("nan") control_lo = np.quantile(control_norm, 0.25) if control_norm.size else float("nan") push_hi = np.quantile(abs_push, 0.75) if abs_push.size else float("nan") push_lo = np.quantile(abs_push, 0.25) if abs_push.size else float("nan") cross_threshold = float(np.quantile(abs_slice_cross_energy, cross_threshold_quantile)) if abs_slice_cross_energy.size else float("nan") threshold_cross = abs_slice_cross_energy >= cross_threshold if np.isfinite(cross_threshold) else np.zeros_like(abs_slice_cross_energy, dtype=bool) high_control_mask = control_norm >= control_hi if np.isfinite(control_hi) else np.zeros_like(control_norm, dtype=bool) low_control_mask = control_norm <= control_lo if np.isfinite(control_lo) else np.zeros_like(control_norm, dtype=bool) high_push_mask = abs_push >= push_hi if np.isfinite(push_hi) else np.zeros_like(abs_push, dtype=bool) low_push_mask = abs_push <= push_lo if np.isfinite(push_lo) else np.zeros_like(abs_push, dtype=bool) sweep_lift_auc = float("nan") sweep_binary_corr_auc = float("nan") sweep_best_threshold_quantile = float("nan") sweep_best_lift = float("nan") sweep_best_binary_corr = float("nan") return PairSummary( q_index=pair.q_index, p_index=pair.p_index, kind=pair.kind, score=pair.score, num_segments=int(abs_delta_h.size), mean_abs_delta_h=float(np.mean(abs_delta_h)), mean_control_norm=float(np.mean(control_norm)), mean_abs_control_energy_push=float(np.mean(abs_push)), mean_cross_ratio=float(np.mean(cross_ratio)), mean_abs_slice_cross_energy=float(np.mean(abs_slice_cross_energy)), cross_threshold=cross_threshold, threshold_cross_rate=float(np.mean(threshold_cross.astype(np.float64))), corr_abs_delta_h_control_norm=_safe_pearson(abs_delta_h, control_norm), corr_abs_delta_h_abs_control_energy_push=_safe_pearson(abs_delta_h, abs_push), corr_cross_ratio_control_norm=_safe_pearson(cross_ratio, control_norm), corr_cross_ratio_abs_control_energy_push=_safe_pearson(cross_ratio, abs_push), corr_signed_delta_h_control_energy_push=_safe_pearson(signed_delta_h, push), corr_abs_slice_cross_energy_abs_control_energy_push=_safe_pearson(abs_slice_cross_energy, abs_push), corr_binary_cross_abs_control_energy_push=_safe_pearson(threshold_cross.astype(np.float64), abs_push), mean_cross_ratio_high_control=_quantile_mean(cross_ratio, high_control_mask), mean_cross_ratio_low_control=_quantile_mean(cross_ratio, low_control_mask), mean_abs_delta_h_high_control=_quantile_mean(abs_delta_h, high_control_mask), mean_abs_delta_h_low_control=_quantile_mean(abs_delta_h, low_control_mask), mean_cross_ratio_high_push=_quantile_mean(cross_ratio, high_push_mask), mean_cross_ratio_low_push=_quantile_mean(cross_ratio, low_push_mask), mean_abs_delta_h_high_push=_quantile_mean(abs_delta_h, high_push_mask), mean_abs_delta_h_low_push=_quantile_mean(abs_delta_h, low_push_mask), threshold_cross_rate_high_push=_quantile_mean(threshold_cross.astype(np.float64), high_push_mask), threshold_cross_rate_low_push=_quantile_mean(threshold_cross.astype(np.float64), low_push_mask), sweep_lift_auc=sweep_lift_auc, sweep_binary_corr_auc=sweep_binary_corr_auc, sweep_best_threshold_quantile=sweep_best_threshold_quantile, sweep_best_lift=sweep_best_lift, sweep_best_binary_corr=sweep_best_binary_corr, ) def compute_threshold_sweep( steps: dict[str, np.ndarray], threshold_quantiles: list[float], high_push_quantile: float = 0.75, low_push_quantile: float = 0.25, push_bins: int = 8, ) -> tuple[list[dict[str, float]], np.ndarray, np.ndarray]: abs_push = steps["abs_control_energy_push"] abs_slice_cross_energy = steps["abs_slice_cross_energy"] high_push_mask = _quantile_mask(abs_push, high_push_quantile, upper=True) low_push_mask = _quantile_mask(abs_push, low_push_quantile, upper=False) x_min = float(np.min(abs_push)) if abs_push.size else 0.0 x_max = float(np.max(abs_push)) if abs_push.size else 1.0 if x_max <= x_min: edges = np.linspace(0.0, 1.0, push_bins + 1) else: edges = np.linspace(x_min, x_max, push_bins + 1) centers = 0.5 * (edges[:-1] + edges[1:]) heatmap_rows: list[np.ndarray] = [] sweep_rows: list[dict[str, float]] = [] for quantile in threshold_quantiles: threshold = float(np.quantile(abs_slice_cross_energy, quantile)) crossing = abs_slice_cross_energy >= threshold high_rate = _quantile_mean(crossing.astype(np.float64), high_push_mask) low_rate = _quantile_mean(crossing.astype(np.float64), low_push_mask) binary_corr = _safe_pearson(crossing.astype(np.float64), abs_push) amplitude_corr = _safe_pearson( np.where(crossing, abs_slice_cross_energy, 0.0), abs_push, ) row = { "threshold_quantile": float(quantile), "cross_threshold": threshold, "overall_cross_rate": float(np.mean(crossing.astype(np.float64))), "high_push_cross_rate": high_rate, "low_push_cross_rate": low_rate, "push_cross_rate_gap": float(high_rate - low_rate), "binary_cross_push_corr": binary_corr, "amplitude_cross_push_corr": amplitude_corr, } sweep_rows.append(row) bin_rates = [] for left, right in zip(edges[:-1], edges[1:]): if right == edges[-1]: mask = (abs_push >= left) & (abs_push <= right) else: mask = (abs_push >= left) & (abs_push < right) bin_rates.append(float(np.mean(crossing[mask].astype(np.float64))) if np.any(mask) else float("nan")) heatmap_rows.append(np.asarray(bin_rates, dtype=np.float64)) return sweep_rows, centers, np.asarray(heatmap_rows, dtype=np.float64) def attach_sweep_metrics(summary: PairSummary, sweep_rows: list[dict[str, float]]) -> PairSummary: if not sweep_rows: return summary lifts = np.asarray([row["push_cross_rate_gap"] for row in sweep_rows], dtype=np.float64) corrs = np.asarray([row["binary_cross_push_corr"] for row in sweep_rows], dtype=np.float64) best_index = int(np.nanargmax(lifts)) best_row = sweep_rows[best_index] return replace( summary, sweep_lift_auc=float(np.nanmean(lifts)), sweep_binary_corr_auc=float(np.nanmean(corrs)), sweep_best_threshold_quantile=float(best_row["threshold_quantile"]), sweep_best_lift=float(best_row["push_cross_rate_gap"]), sweep_best_binary_corr=float(best_row["binary_cross_push_corr"]), ) def _draw_background(axis: plt.Axes, slice_eval: dict[str, np.ndarray], pair: PairScore, title: str) -> None: q_axis = slice_eval["q_axis"] p_axis = slice_eval["p_axis"] axis.contourf(q_axis, p_axis, slice_eval["H"], levels=28, cmap="magma", alpha=0.95) axis.contour( q_axis, p_axis, slice_eval["H"], levels=10, colors="white", linewidths=0.8, linestyles="--", alpha=0.48, ) axis.streamplot( q_axis, p_axis, slice_eval["U"], slice_eval["V"], density=1.05, color="#203a8a", linewidth=1.15, arrowsize=1.35, maxlength=3.5, ) axis.set_xlabel(f"q[{pair.q_index}]", fontsize=PAPER_LABEL_FONTSIZE) axis.set_ylabel(f"p[{pair.p_index}]", fontsize=PAPER_LABEL_FONTSIZE) axis.set_title(title, loc="left", fontsize=PAPER_TITLE_FONTSIZE, pad=12) axis.tick_params(axis="both", labelsize=PAPER_TICK_FONTSIZE, width=PAPER_SPINE_WIDTH, length=7, pad=8) axis.spines["left"].set_linewidth(PAPER_SPINE_WIDTH) axis.spines["bottom"].set_linewidth(PAPER_SPINE_WIDTH) axis.spines["top"].set_visible(False) axis.spines["right"].set_visible(False) def _line_collection( axis: plt.Axes, segments: np.ndarray, values: np.ndarray, cmap: str, norm: Normalize, linewidth: float = 2.4, ) -> LineCollection: collection = LineCollection( segments, cmap=cmap, norm=norm, linewidths=linewidth, alpha=0.98, capstyle="round", joinstyle="round", zorder=5, ) collection.set_array(values) axis.add_collection(collection) return collection def _scatter_points(axis: plt.Axes, steps: dict[str, np.ndarray]) -> None: axis.scatter(steps["x0"][0], steps["y0"][0], color="#7dd3fc", s=60, zorder=6, label="start") axis.scatter(steps["x1"][-1], steps["y1"][-1], color="#38bdf8", s=60, zorder=6, label="end") def _style_colorbar(cbar: plt.colorbar, label: str) -> None: cbar.set_label(label, fontsize=PAPER_CBAR_FONTSIZE, labelpad=14) cbar.ax.tick_params(labelsize=PAPER_TICK_FONTSIZE, width=PAPER_SPINE_WIDTH, length=6) def _style_axis(axis: plt.Axes) -> None: axis.tick_params(axis="both", labelsize=PAPER_TICK_FONTSIZE, width=PAPER_SPINE_WIDTH, length=7, pad=8) axis.spines["left"].set_linewidth(PAPER_SPINE_WIDTH) axis.spines["bottom"].set_linewidth(PAPER_SPINE_WIDTH) axis.spines["top"].set_visible(False) axis.spines["right"].set_visible(False) def draw_pair_overlay_and_scatter( output_path: Path, slice_eval: dict[str, np.ndarray], pair: PairScore, steps: dict[str, np.ndarray], summary: PairSummary, scatter_max_points: int, ) -> None: delta_h = steps["delta_h"] push = steps["control_energy_push"] cross_ratio = steps["cross_ratio"] control_norm = steps["control_norm"] abs_push = steps["abs_control_energy_push"] abs_delta_h = steps["abs_delta_h"] segments = steps["segments"] delta_scale = max(1e-8, float(np.max(np.abs(delta_h)))) push_scale = max(1e-8, float(np.max(np.abs(push)))) delta_norm = TwoSlopeNorm(vmin=-delta_scale, vcenter=0.0, vmax=delta_scale) push_norm = TwoSlopeNorm(vmin=-push_scale, vcenter=0.0, vmax=push_scale) cross_norm = Normalize(vmin=0.0, vmax=1.0) fig, axes = plt.subplots(2, 3, figsize=(15.6, 8.8), dpi=220) titles = [ r"Trajectory colored by $\Delta H_t$", r"Trajectory colored by $(\nabla_p H_t)\cdot control_t$", r"Trajectory colored by crossing ratio", ] overlay_specs = [ (delta_h, "coolwarm", delta_norm), (push, "coolwarm", push_norm), (cross_ratio, "viridis", cross_norm), ] cbar_labels = [ r"$\Delta H_t = H_{t+1} - H_t$", r"$(\nabla_p H_t)\cdot control_t$", r"$|u_t \cdot n_t| / (||u_t||\,||n_t||)$", ] for axis, title, (values, cmap, norm), cbar_label in zip(axes[0], titles, overlay_specs, cbar_labels): _draw_background(axis, slice_eval, pair, title) collection = _line_collection(axis, segments, values, cmap=cmap, norm=norm) _scatter_points(axis, steps) cbar = fig.colorbar(collection, ax=axis, fraction=0.046, pad=0.03) cbar.set_label(cbar_label) rng = np.random.default_rng(42) sample_size = min(int(scatter_max_points), int(abs_delta_h.size)) sample_idx = np.arange(abs_delta_h.size) if sample_size < abs_delta_h.size: sample_idx = rng.choice(abs_delta_h.size, size=sample_size, replace=False) scatter_specs = [ ( control_norm[sample_idx], cross_ratio[sample_idx], abs_push[sample_idx], r"$||control_t||_2$", "cross ratio", f"cross vs control\nr={summary.corr_cross_ratio_control_norm:.3f}", ), ( abs_push[sample_idx], cross_ratio[sample_idx], control_norm[sample_idx], r"$|(\nabla_p H_t)\cdot control_t|$", "cross ratio", f"cross vs |push|\nr={summary.corr_cross_ratio_abs_control_energy_push:.3f}", ), ( abs_push[sample_idx], abs_delta_h[sample_idx], control_norm[sample_idx], r"$|(\nabla_p H_t)\cdot control_t|$", r"$|\Delta H_t|$", f"|ΔH| vs |push|\nr={summary.corr_abs_delta_h_abs_control_energy_push:.3f}", ), ] for axis, (xs, ys, color_values, xlabel, ylabel, title) in zip(axes[1], scatter_specs): scatter = axis.scatter(xs, ys, c=color_values, cmap="plasma", s=18, alpha=0.72, edgecolors="none") axis.set_xlabel(xlabel) axis.set_ylabel(ylabel) axis.set_title(title, loc="left", fontsize=10) axis.spines["top"].set_visible(False) axis.spines["right"].set_visible(False) axis.grid(True, alpha=0.25) cbar = fig.colorbar(scatter, ax=axis, fraction=0.046, pad=0.03) cbar.set_label(r"$||control_t||_2$" if xlabel != r"$||control_t||_2$" else r"$|(\nabla_p H_t)\cdot control_t|$") fig.suptitle( f"{pair.kind} pair q[{pair.q_index}] / p[{pair.p_index}] score={pair.score:.3f}\n" f"corr(sign ΔH, push)={summary.corr_signed_delta_h_control_energy_push:.3f}", x=0.055, y=0.995, ha="left", fontsize=14, ) fig.tight_layout(rect=[0, 0, 1, 0.965]) fig.savefig(output_path, bbox_inches="tight") plt.close(fig) def draw_pair_triptych( output_path: Path, slice_eval: dict[str, np.ndarray], pair: PairScore, steps: dict[str, np.ndarray], ) -> None: delta_h = steps["delta_h"] push = steps["control_energy_push"] cross_ratio = steps["cross_ratio"] segments = steps["segments"] delta_scale = max(1e-8, float(np.max(np.abs(delta_h)))) push_scale = max(1e-8, float(np.max(np.abs(push)))) fig, axes = plt.subplots(1, 3, figsize=(16.0, 4.9), dpi=220) configs = [ (delta_h, "coolwarm", TwoSlopeNorm(vmin=-delta_scale, vcenter=0.0, vmax=delta_scale), r"$\Delta H_t$"), (push, "coolwarm", TwoSlopeNorm(vmin=-push_scale, vcenter=0.0, vmax=push_scale), r"$(\nabla_p H_t)\cdot control_t$"), (cross_ratio, "viridis", Normalize(vmin=0.0, vmax=1.0), "cross ratio"), ] for axis, (values, cmap, norm, label) in zip(axes, configs): _draw_background(axis, slice_eval, pair, label) collection = _line_collection(axis, segments, values, cmap=cmap, norm=norm) _scatter_points(axis, steps) cbar = fig.colorbar(collection, ax=axis, fraction=0.046, pad=0.03) cbar.set_label(label) fig.tight_layout() fig.savefig(output_path, bbox_inches="tight") plt.close(fig) def draw_pair_push_flow( output_path: Path, slice_eval: dict[str, np.ndarray], pair: PairScore, steps: dict[str, np.ndarray], summary: PairSummary | None = None, ) -> None: push = steps["control_energy_push"] segments = steps["segments"] push_scale = max(1e-8, float(np.max(np.abs(push)))) push_norm = TwoSlopeNorm(vmin=-push_scale, vcenter=0.0, vmax=push_scale) fig, axis = plt.subplots(1, 1, figsize=(8.4, 6.8), dpi=260) title = r"Trajectory colored by $(\nabla_p H_t)\cdot control_t$" if summary is not None: title += ( f"\n|ΔH|-|push| r={summary.corr_abs_delta_h_abs_control_energy_push:.3f}, " f"lift_auc={summary.sweep_lift_auc:.3f}" ) _draw_background(axis, slice_eval, pair, title) collection = _line_collection(axis, segments, push, cmap="coolwarm", norm=push_norm, linewidth=4.0) _scatter_points(axis, steps) cbar = fig.colorbar(collection, ax=axis, fraction=0.046, pad=0.04) _style_colorbar(cbar, r"$(\nabla_p H_t)\cdot control_t$") fig.tight_layout(pad=0.8) fig.savefig(output_path, bbox_inches="tight", pad_inches=0.05) plt.close(fig) def _binned_mean_curve(x: np.ndarray, y: np.ndarray, bins: int = 8) -> tuple[np.ndarray, np.ndarray]: if x.size == 0: return np.asarray([]), np.asarray([]) x_min = float(np.min(x)) x_max = float(np.max(x)) if not np.isfinite(x_min) or not np.isfinite(x_max) or x_max <= x_min: return np.asarray([]), np.asarray([]) edges = np.linspace(x_min, x_max, bins + 1) centers = 0.5 * (edges[:-1] + edges[1:]) means = [] for left, right in zip(edges[:-1], edges[1:]): if right == edges[-1]: mask = (x >= left) & (x <= right) else: mask = (x >= left) & (x < right) means.append(float(np.mean(y[mask])) if np.any(mask) else float("nan")) return centers, np.asarray(means, dtype=np.float64) def draw_pair_thresholded( output_path: Path, pair: PairScore, steps: dict[str, np.ndarray], summary: PairSummary, scatter_max_points: int, ) -> None: abs_push = steps["abs_control_energy_push"] abs_slice_cross_energy = steps["abs_slice_cross_energy"] step_norm = steps["step_norm"] threshold_cross = (abs_slice_cross_energy >= summary.cross_threshold).astype(np.float64) rng = np.random.default_rng(42) sample_size = min(int(scatter_max_points), int(abs_push.size)) sample_idx = np.arange(abs_push.size) if sample_size < abs_push.size: sample_idx = rng.choice(abs_push.size, size=sample_size, replace=False) jitter = rng.normal(loc=0.0, scale=0.035, size=sample_idx.size) x_bins, y_bins = _binned_mean_curve(abs_push, threshold_cross, bins=8) x_bins_energy, y_bins_energy = _binned_mean_curve(abs_push, abs_slice_cross_energy, bins=8) fig, axes = plt.subplots(1, 3, figsize=(14.8, 4.2), dpi=220) scatter0 = axes[0].scatter( abs_push[sample_idx], abs_slice_cross_energy[sample_idx], c=step_norm[sample_idx], cmap="plasma", s=20, alpha=0.76, edgecolors="none", ) axes[0].axhline(summary.cross_threshold, color="#0f172a", linewidth=1.2, linestyle="--", alpha=0.85) axes[0].set_xlabel(r"$|(\nabla_p H_t)\cdot control_t|$") axes[0].set_ylabel(r"$|(\nabla_{slice} H_t)\cdot \Delta x_t|$") axes[0].set_title( f"slice cross amplitude vs |push|\nr={summary.corr_abs_slice_cross_energy_abs_control_energy_push:.3f}", loc="left", fontsize=10, ) axes[0].grid(True, alpha=0.25) axes[0].spines["top"].set_visible(False) axes[0].spines["right"].set_visible(False) cbar0 = fig.colorbar(scatter0, ax=axes[0], fraction=0.046, pad=0.03) cbar0.set_label(r"$||\Delta x_t||_2$") scatter1 = axes[1].scatter( abs_push[sample_idx], threshold_cross[sample_idx] + jitter, c=step_norm[sample_idx], cmap="viridis", s=20, alpha=0.78, edgecolors="none", ) axes[1].set_xlabel(r"$|(\nabla_p H_t)\cdot control_t|$") axes[1].set_ylabel("thresholded crossing") axes[1].set_yticks([0.0, 1.0]) axes[1].set_yticklabels(["no", "yes"]) axes[1].set_title( f"binary crossing vs |push|\nr={summary.corr_binary_cross_abs_control_energy_push:.3f}", loc="left", fontsize=10, ) axes[1].grid(True, alpha=0.25) axes[1].spines["top"].set_visible(False) axes[1].spines["right"].set_visible(False) cbar1 = fig.colorbar(scatter1, ax=axes[1], fraction=0.046, pad=0.03) cbar1.set_label(r"$||\Delta x_t||_2$") axes[2].plot(x_bins, y_bins, color="#dc2626", linewidth=2.2, marker="o", label="crossing rate") axes[2].plot(x_bins_energy, y_bins_energy, color="#2563eb", linewidth=2.0, marker="s", label="mean |slice cross|") axes[2].axhline(summary.threshold_cross_rate, color="#64748b", linewidth=1.0, linestyle="--", alpha=0.7) axes[2].set_xlabel(r"$|(\nabla_p H_t)\cdot control_t|$ bin center") axes[2].set_ylabel("rate / mean") axes[2].set_title( f"binned trend\nhigh push rate={summary.threshold_cross_rate_high_push:.3f}, low push rate={summary.threshold_cross_rate_low_push:.3f}", loc="left", fontsize=10, ) axes[2].grid(True, alpha=0.25) axes[2].legend(frameon=False, loc="best", fontsize=8) axes[2].spines["top"].set_visible(False) axes[2].spines["right"].set_visible(False) fig.suptitle( f"{pair.kind} pair q[{pair.q_index}] / p[{pair.p_index}] thresholded slice crossing\n" f"cross threshold={summary.cross_threshold:.4g}", x=0.055, y=0.995, ha="left", fontsize=13, ) fig.tight_layout(rect=[0, 0, 1, 0.95]) fig.savefig(output_path, bbox_inches="tight") plt.close(fig) def draw_pair_threshold_sweep( output_path: Path, pair: PairScore, summary: PairSummary, sweep_rows: list[dict[str, float]], push_bin_centers: np.ndarray, sweep_heatmap: np.ndarray, ) -> None: quantiles = np.asarray([row["threshold_quantile"] for row in sweep_rows], dtype=np.float64) high_rates = np.asarray([row["high_push_cross_rate"] for row in sweep_rows], dtype=np.float64) low_rates = np.asarray([row["low_push_cross_rate"] for row in sweep_rows], dtype=np.float64) gaps = np.asarray([row["push_cross_rate_gap"] for row in sweep_rows], dtype=np.float64) binary_corrs = np.asarray([row["binary_cross_push_corr"] for row in sweep_rows], dtype=np.float64) fig, axes = plt.subplots(1, 3, figsize=(15.2, 4.3), dpi=220) axes[0].plot(quantiles, high_rates, color="#dc2626", linewidth=2.2, marker="o", label="high push") axes[0].plot(quantiles, low_rates, color="#2563eb", linewidth=2.2, marker="s", label="low push") axes[0].plot(quantiles, gaps, color="#0f172a", linewidth=2.0, marker="^", label="gap") axes[0].set_xlabel("cross threshold quantile") axes[0].set_ylabel("crossing rate / gap") axes[0].set_title( f"threshold sweep\nlift_auc={summary.sweep_lift_auc:.3f}", loc="left", fontsize=10, ) axes[0].grid(True, alpha=0.25) axes[0].legend(frameon=False, loc="best", fontsize=8) axes[0].spines["top"].set_visible(False) axes[0].spines["right"].set_visible(False) heat = axes[1].imshow( sweep_heatmap, aspect="auto", origin="lower", cmap="YlGnBu", extent=[float(push_bin_centers[0]), float(push_bin_centers[-1]), float(quantiles[0]), float(quantiles[-1])], vmin=0.0, vmax=1.0, ) axes[1].set_xlabel(r"$|(\nabla_p H_t)\cdot control_t|$ bin center") axes[1].set_ylabel("cross threshold quantile") axes[1].set_title("crossing probability heatmap", loc="left", fontsize=10) axes[1].spines["top"].set_visible(False) axes[1].spines["right"].set_visible(False) cbar = fig.colorbar(heat, ax=axes[1], fraction=0.046, pad=0.03) cbar.set_label("P(crossing)") axes[2].plot(quantiles, binary_corrs, color="#7c3aed", linewidth=2.2, marker="o", label="binary corr") axes[2].axhline(summary.sweep_binary_corr_auc, color="#64748b", linewidth=1.1, linestyle="--", label="corr auc") axes[2].set_xlabel("cross threshold quantile") axes[2].set_ylabel(r"corr(binary crossing, $|push|$)") axes[2].set_title( f"correlation sweep\nbest_q={summary.sweep_best_threshold_quantile:.2f}, best_lift={summary.sweep_best_lift:.3f}", loc="left", fontsize=10, ) axes[2].grid(True, alpha=0.25) axes[2].legend(frameon=False, loc="best", fontsize=8) axes[2].spines["top"].set_visible(False) axes[2].spines["right"].set_visible(False) fig.suptitle( f"{pair.kind} pair q[{pair.q_index}] / p[{pair.p_index}] threshold sweep view", x=0.055, y=0.995, ha="left", fontsize=13, ) fig.tight_layout(rect=[0, 0, 1, 0.95]) fig.savefig(output_path, bbox_inches="tight") plt.close(fig) def draw_pair_push_crossing_rate_curve( output_path: Path, pair: PairScore, summary: PairSummary, sweep_rows: list[dict[str, float]], ) -> None: quantiles = np.asarray([row["threshold_quantile"] for row in sweep_rows], dtype=np.float64) high_rates = np.asarray([row["high_push_cross_rate"] for row in sweep_rows], dtype=np.float64) low_rates = np.asarray([row["low_push_cross_rate"] for row in sweep_rows], dtype=np.float64) fig, axis = plt.subplots(1, 1, figsize=(7.8, 5.8), dpi=260) axis.plot(quantiles, high_rates, color="#dc2626", linewidth=3.0, marker="o", markersize=8.0, label="High |push|") axis.plot(quantiles, low_rates, color="#2563eb", linewidth=3.0, marker="s", markersize=8.0, label="Low |push|") axis.fill_between(quantiles, low_rates, high_rates, color="#fca5a5", alpha=0.24) axis.set_xlabel("Crossing Threshold Quantile", fontsize=PAPER_LABEL_FONTSIZE) axis.set_ylabel("Crossing Rate", fontsize=PAPER_LABEL_FONTSIZE) axis.set_title( f"{pair.kind} q[{pair.q_index}] / p[{pair.p_index}]\n" f"lift_auc={summary.sweep_lift_auc:.3f}, best_lift={summary.sweep_best_lift:.3f}", loc="left", fontsize=PAPER_TITLE_FONTSIZE, pad=12, ) axis.set_xlim(float(np.min(quantiles)) - 0.01, float(np.max(quantiles)) + 0.01) axis.set_ylim(-0.02, 1.02) axis.set_xticks(quantiles) axis.grid(True, alpha=0.22, linewidth=0.8) _style_axis(axis) axis.legend(frameon=False, loc="best", fontsize=PAPER_LEGEND_FONTSIZE + 3, handlelength=2.8) fig.tight_layout(pad=0.8) fig.savefig(output_path, bbox_inches="tight", pad_inches=0.05) plt.close(fig) def _pair_label(pair: PairScore | PairSummary) -> str: return f"{pair.kind} q[{pair.q_index}] / p[{pair.p_index}]" def _summary_matrix( summaries: list[PairSummary], value_key: str, q_dim: int, p_dim: int, ) -> np.ndarray: matrix = np.full((q_dim, p_dim), np.nan, dtype=np.float64) for summary in summaries: matrix[summary.q_index, summary.p_index] = float(getattr(summary, value_key)) return matrix def draw_all_pair_rankings( output_path: Path, summaries: list[PairSummary], q_dim: int, p_dim: int, top_k: int, ) -> None: ordered_by_lift = sorted(summaries, key=lambda row: row.sweep_lift_auc, reverse=True) ordered_by_corr = sorted(summaries, key=lambda row: row.sweep_binary_corr_auc, reverse=True) top_lift = ordered_by_lift[: max(1, int(top_k))] top_corr = ordered_by_corr[: max(1, int(top_k))] fig, axes = plt.subplots(2, 2, figsize=(14.6, 10.0), dpi=220) lift_labels = [_pair_label(row) for row in reversed(top_lift)] lift_values = [row.sweep_lift_auc for row in reversed(top_lift)] axes[0, 0].barh(lift_labels, lift_values, color="#dc2626", alpha=0.88) axes[0, 0].set_title("Top Pairs By Sweep Lift AUC", loc="left") axes[0, 0].set_xlabel("sweep_lift_auc") axes[0, 0].grid(True, axis="x", alpha=0.22) axes[0, 0].spines["top"].set_visible(False) axes[0, 0].spines["right"].set_visible(False) corr_labels = [_pair_label(row) for row in reversed(top_corr)] corr_values = [row.sweep_binary_corr_auc for row in reversed(top_corr)] axes[0, 1].barh(corr_labels, corr_values, color="#7c3aed", alpha=0.88) axes[0, 1].set_title("Top Pairs By Sweep Binary Corr AUC", loc="left") axes[0, 1].set_xlabel("sweep_binary_corr_auc") axes[0, 1].grid(True, axis="x", alpha=0.22) axes[0, 1].spines["top"].set_visible(False) axes[0, 1].spines["right"].set_visible(False) lift_matrix = _summary_matrix(summaries, "sweep_lift_auc", q_dim=q_dim, p_dim=p_dim) corr_matrix = _summary_matrix(summaries, "sweep_binary_corr_auc", q_dim=q_dim, p_dim=p_dim) im0 = axes[1, 0].imshow(lift_matrix, cmap="YlOrRd", aspect="auto") axes[1, 0].set_title("Sweep Lift AUC Matrix", loc="left") axes[1, 0].set_xlabel("p index") axes[1, 0].set_ylabel("q index") axes[1, 0].set_xticks(np.arange(p_dim)) axes[1, 0].set_yticks(np.arange(q_dim)) cbar0 = fig.colorbar(im0, ax=axes[1, 0], fraction=0.046, pad=0.03) cbar0.set_label("sweep_lift_auc") im1 = axes[1, 1].imshow(corr_matrix, cmap="BuPu", aspect="auto") axes[1, 1].set_title("Sweep Binary Corr AUC Matrix", loc="left") axes[1, 1].set_xlabel("p index") axes[1, 1].set_ylabel("q index") axes[1, 1].set_xticks(np.arange(p_dim)) axes[1, 1].set_yticks(np.arange(q_dim)) cbar1 = fig.colorbar(im1, ax=axes[1, 1], fraction=0.046, pad=0.03) cbar1.set_label("sweep_binary_corr_auc") fig.suptitle("All Pair Rankings And Metric Maps", x=0.055, y=0.995, ha="left", fontsize=15) fig.tight_layout(rect=[0, 0, 1, 0.97]) fig.savefig(output_path, bbox_inches="tight") plt.close(fig) def draw_all_pair_focus_heatmap( output_path: Path, summaries: list[PairSummary], q_dim: int, p_dim: int, ) -> None: best_lift_matrix = _summary_matrix(summaries, "sweep_best_lift", q_dim=q_dim, p_dim=p_dim) best_q_matrix = _summary_matrix(summaries, "sweep_best_threshold_quantile", q_dim=q_dim, p_dim=p_dim) fig, axes = plt.subplots(1, 2, figsize=(12.2, 4.8), dpi=220) im0 = axes[0].imshow(best_lift_matrix, cmap="magma", aspect="auto") axes[0].set_title("Best Threshold Lift", loc="left") axes[0].set_xlabel("p index") axes[0].set_ylabel("q index") axes[0].set_xticks(np.arange(p_dim)) axes[0].set_yticks(np.arange(q_dim)) fig.colorbar(im0, ax=axes[0], fraction=0.046, pad=0.03).set_label("best_lift") im1 = axes[1].imshow(best_q_matrix, cmap="viridis", aspect="auto", vmin=0.4, vmax=0.9) axes[1].set_title("Best Threshold Quantile", loc="left") axes[1].set_xlabel("p index") axes[1].set_ylabel("q index") axes[1].set_xticks(np.arange(p_dim)) axes[1].set_yticks(np.arange(q_dim)) fig.colorbar(im1, ax=axes[1], fraction=0.046, pad=0.03).set_label("best threshold q") fig.tight_layout() fig.savefig(output_path, bbox_inches="tight") plt.close(fig) def draw_focus_story_sheet( output_path: Path, slice_eval: dict[str, np.ndarray], pair: PairScore, steps: dict[str, np.ndarray], summary: PairSummary, sweep_rows: list[dict[str, float]], ) -> None: delta_h = steps["delta_h"] push = steps["control_energy_push"] abs_push = steps["abs_control_energy_push"] abs_delta_h = steps["abs_delta_h"] segments = steps["segments"] delta_scale = max(1e-8, float(np.max(np.abs(delta_h)))) push_scale = max(1e-8, float(np.max(np.abs(push)))) delta_norm = TwoSlopeNorm(vmin=-delta_scale, vcenter=0.0, vmax=delta_scale) push_norm = TwoSlopeNorm(vmin=-push_scale, vcenter=0.0, vmax=push_scale) quantiles = np.asarray([row["threshold_quantile"] for row in sweep_rows], dtype=np.float64) high_rates = np.asarray([row["high_push_cross_rate"] for row in sweep_rows], dtype=np.float64) low_rates = np.asarray([row["low_push_cross_rate"] for row in sweep_rows], dtype=np.float64) gaps = np.asarray([row["push_cross_rate_gap"] for row in sweep_rows], dtype=np.float64) fig = plt.figure(figsize=(14.8, 8.2), dpi=220) gs = fig.add_gridspec(2, 3, width_ratios=[1.2, 1.0, 0.95], height_ratios=[1.0, 1.0]) ax0 = fig.add_subplot(gs[:, 0]) _draw_background(ax0, slice_eval, pair, r"Trajectory colored by $(\nabla_p H_t)\cdot control_t$") coll0 = _line_collection(ax0, segments, push, cmap="coolwarm", norm=push_norm, linewidth=2.8) _scatter_points(ax0, steps) fig.colorbar(coll0, ax=ax0, fraction=0.046, pad=0.03).set_label(r"$(\nabla_p H_t)\cdot control_t$") ax1 = fig.add_subplot(gs[0, 1]) scatter1 = ax1.scatter(abs_push, abs_delta_h, c=steps["step_norm"], cmap="plasma", s=22, alpha=0.75, edgecolors="none") ax1.set_xlabel(r"$|(\nabla_p H_t)\cdot control_t|$") ax1.set_ylabel(r"$|\Delta H_t|$") ax1.set_title(f"Energy response\nr={summary.corr_abs_delta_h_abs_control_energy_push:.3f}", loc="left") ax1.grid(True, alpha=0.25) ax1.spines["top"].set_visible(False) ax1.spines["right"].set_visible(False) fig.colorbar(scatter1, ax=ax1, fraction=0.046, pad=0.03).set_label(r"$||\Delta x_t||_2$") ax2 = fig.add_subplot(gs[1, 1]) ax2.plot(quantiles, high_rates, color="#dc2626", linewidth=2.2, marker="o", label="high push") ax2.plot(quantiles, low_rates, color="#2563eb", linewidth=2.2, marker="s", label="low push") ax2.plot(quantiles, gaps, color="#111827", linewidth=2.0, marker="^", label="gap") ax2.set_xlabel("cross threshold quantile") ax2.set_ylabel("crossing rate / gap") ax2.set_title( f"Threshold sweep\nlift_auc={summary.sweep_lift_auc:.3f}, best_lift={summary.sweep_best_lift:.3f}", loc="left", ) ax2.grid(True, alpha=0.25) ax2.legend(frameon=False, loc="best", fontsize=8) ax2.spines["top"].set_visible(False) ax2.spines["right"].set_visible(False) ax3 = fig.add_subplot(gs[:, 2]) ax3.axis("off") lines = [ f"Focus Pair: {_pair_label(pair)}", f"score = {pair.score:.3f}", "", "Key numbers", f"corr(sign ΔH, push) = {summary.corr_signed_delta_h_control_energy_push:.3f}", f"corr(|ΔH|, |push|) = {summary.corr_abs_delta_h_abs_control_energy_push:.3f}", f"corr(binary crossing, |push|) = {summary.corr_binary_cross_abs_control_energy_push:.3f}", f"sweep_lift_auc = {summary.sweep_lift_auc:.3f}", f"sweep_binary_corr_auc = {summary.sweep_binary_corr_auc:.3f}", f"best threshold q = {summary.sweep_best_threshold_quantile:.2f}", f"best lift = {summary.sweep_best_lift:.3f}", "", "Interpretation", "Large |push| segments tend to coincide", "with larger |ΔH| and, after thresholding,", "with a higher chance of visible crossing.", "", "This supports:", "control projected onto ∇_p H explains", "energy-changing events better than", "plain control magnitude alone.", ] ax3.text( 0.0, 1.0, "\n".join(lines), va="top", ha="left", fontsize=11, family="monospace", color="#0f172a", ) fig.suptitle("Focused Hamiltonian Control-Crossing Story", x=0.055, y=0.995, ha="left", fontsize=15) fig.tight_layout(rect=[0, 0, 1, 0.97]) fig.savefig(output_path, bbox_inches="tight") plt.close(fig) def write_csv(output_path: Path, rows: list[dict[str, object]]) -> None: if not rows: return fieldnames = list(rows[0].keys()) with output_path.open("w", encoding="utf-8", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=fieldnames) writer.writeheader() writer.writerows(rows) def main() -> int: args = parse_args() checkpoint_path = Path(args.checkpoint).expanduser().resolve() trace_path = Path(args.trace_path).expanduser().resolve() output_dir = Path(args.output_dir).expanduser().resolve() output_dir.mkdir(parents=True, exist_ok=True) threshold_quantiles = _parse_float_list(args.cross_threshold_sweep) device = torch.device(args.device) world_model = _load_checkpoint_world_model(checkpoint_path, device) trace = _load_trace(trace_path) valid_mask = np.asarray(trace["valid_mask"], dtype=bool) q_flat = _flatten_valid_vectors(trace["q"], valid_mask) p_flat = _flatten_valid_vectors(trace["p"], valid_mask) q_reference = np.median(q_flat, axis=0).astype(np.float32) p_reference = np.median(p_flat, axis=0).astype(np.float32) q_dim = int(np.asarray(trace["q_dim"]).item()) p_dim = int(np.asarray(trace["p_dim"]).item()) all_scores = score_pairs(trace) selected_pairs = select_pairs(all_scores, same_k=args.same_k, cross_k=args.cross_k) if not selected_pairs: raise RuntimeError("No q/p pairs were selected for rendering.") figures_dir = output_dir / "pairs" figures_dir.mkdir(parents=True, exist_ok=True) summary_dir = output_dir / "summary_views" summary_dir.mkdir(parents=True, exist_ok=True) all_pair_summaries: list[PairSummary] = [] pair_summary_map: dict[tuple[int, int], PairSummary] = {} pair_sweep_map: dict[tuple[int, int], list[dict[str, float]]] = {} selection_rows = [asdict(pair) for pair in selected_pairs] write_csv(output_dir / "selected_pairs.csv", selection_rows) write_csv(output_dir / "pair_scores.csv", [asdict(pair) for pair in all_scores]) for pair in all_scores: steps = build_pair_steps(trace, pair) summary = summarize_pair(pair, steps, cross_threshold_quantile=float(args.cross_threshold_quantile)) sweep_rows, _push_bin_centers, _sweep_heatmap = compute_threshold_sweep(steps, threshold_quantiles) summary = attach_sweep_metrics(summary, sweep_rows) all_pair_summaries.append(summary) pair_summary_map[(pair.q_index, pair.p_index)] = summary pair_sweep_map[(pair.q_index, pair.p_index)] = sweep_rows write_csv(output_dir / "all_pair_correlation_summary.csv", [asdict(summary) for summary in all_pair_summaries]) draw_all_pair_rankings( summary_dir / "all_pair_rankings.png", all_pair_summaries, q_dim=q_dim, p_dim=p_dim, top_k=int(args.ranking_top_k), ) draw_all_pair_focus_heatmap( summary_dir / "all_pair_metric_maps.png", all_pair_summaries, q_dim=q_dim, p_dim=p_dim, ) selected_summaries: list[PairSummary] = [] for pair in selected_pairs: slice_eval = evaluate_slice( world_model=world_model, q_reference=q_reference, p_reference=p_reference, pair=pair, q_values=q_flat[:, pair.q_index], p_values=p_flat[:, pair.p_index], grid_size=int(args.grid_size), quantile_lo=float(args.trace_quantile_lo), quantile_hi=float(args.trace_quantile_hi), device=device, ) steps = build_pair_steps(trace, pair) summary = pair_summary_map[(pair.q_index, pair.p_index)] selected_summaries.append(summary) sweep_rows, push_bin_centers, sweep_heatmap = compute_threshold_sweep(steps, threshold_quantiles) step_rows = [] for idx in range(steps["delta_h"].size): step_rows.append( { "episode_id": int(steps["episode_id"][idx]), "step_id": int(steps["step_id"][idx]), "x0": float(steps["x0"][idx]), "y0": float(steps["y0"][idx]), "x1": float(steps["x1"][idx]), "y1": float(steps["y1"][idx]), "delta_h": float(steps["delta_h"][idx]), "abs_delta_h": float(steps["abs_delta_h"][idx]), "control_norm": float(steps["control_norm"][idx]), "control_energy_push": float(steps["control_energy_push"][idx]), "abs_control_energy_push": float(steps["abs_control_energy_push"][idx]), "cross_ratio": float(steps["cross_ratio"][idx]), "abs_slice_cross_energy": float(steps["abs_slice_cross_energy"][idx]), "signed_slice_cross_energy": float(steps["signed_slice_cross_energy"][idx]), "step_norm": float(steps["step_norm"][idx]), } ) pair_stem = f"{pair.kind}_q{pair.q_index}_p{pair.p_index}" write_csv(figures_dir / f"{pair_stem}_steps.csv", step_rows) write_csv(figures_dir / f"{pair_stem}_threshold_sweep.csv", sweep_rows) draw_pair_triptych(figures_dir / f"{pair_stem}_triptych.png", slice_eval, pair, steps) draw_pair_overlay_and_scatter( figures_dir / f"{pair_stem}_analysis.png", slice_eval, pair, steps, summary, scatter_max_points=int(args.scatter_max_points), ) draw_pair_threshold_sweep( figures_dir / f"{pair_stem}_threshold_sweep.png", pair, summary, sweep_rows, push_bin_centers, sweep_heatmap, ) draw_pair_push_crossing_rate_curve( figures_dir / f"{pair_stem}_push_crossing_rate.png", pair, summary, sweep_rows, ) draw_pair_thresholded( figures_dir / f"{pair_stem}_thresholded.png", pair, steps, summary, scatter_max_points=int(args.scatter_max_points), ) write_csv(output_dir / "pair_correlation_summary.csv", [asdict(summary) for summary in selected_summaries]) focus_pair = next( (pair for pair in all_scores if pair.q_index == int(args.focus_q_index) and pair.p_index == int(args.focus_p_index)), None, ) if focus_pair is not None: focus_steps = build_pair_steps(trace, focus_pair) focus_slice = evaluate_slice( world_model=world_model, q_reference=q_reference, p_reference=p_reference, pair=focus_pair, q_values=q_flat[:, focus_pair.q_index], p_values=p_flat[:, focus_pair.p_index], grid_size=int(args.grid_size), quantile_lo=float(args.trace_quantile_lo), quantile_hi=float(args.trace_quantile_hi), device=device, ) draw_focus_story_sheet( summary_dir / f"focus_q{focus_pair.q_index}_p{focus_pair.p_index}_story.png", focus_slice, focus_pair, focus_steps, pair_summary_map[(focus_pair.q_index, focus_pair.p_index)], pair_sweep_map[(focus_pair.q_index, focus_pair.p_index)], ) with (output_dir / "summary.json").open("w", encoding="utf-8") as handle: json.dump( { "checkpoint": str(checkpoint_path), "trace_path": str(trace_path), "task": str(np.asarray(trace["task"]).item()), "seed": int(np.asarray(trace["seed"]).item()), "mode": str(np.asarray(trace["mode"]).item()), "selected_pairs": selection_rows, "selected_pair_summaries": [asdict(summary) for summary in selected_summaries], "all_pair_summary_count": len(all_pair_summaries), "focus_pair": None if focus_pair is None else asdict(focus_pair), }, handle, indent=2, ) for pair in selected_pairs: pair_stem = f"{pair.kind}_q{pair.q_index}_p{pair.p_index}" print(figures_dir / f"{pair_stem}_triptych.png") print(figures_dir / f"{pair_stem}_analysis.png") print(figures_dir / f"{pair_stem}_thresholded.png") print(figures_dir / f"{pair_stem}_threshold_sweep.png") print(figures_dir / f"{pair_stem}_push_crossing_rate.png") print(summary_dir / "all_pair_rankings.png") print(summary_dir / "all_pair_metric_maps.png") if focus_pair is not None: print(summary_dir / f"focus_q{focus_pair.q_index}_p{focus_pair.p_index}_story.png") print(output_dir / "pair_correlation_summary.csv") print(output_dir / "all_pair_correlation_summary.csv") return 0 if __name__ == "__main__": raise SystemExit(main())