Files
DynamisLab/src/drl_pinball/plot_reproduction_summary.py
T
Frank14fandCursor 86157c6f31 feat(drl): publish writing-ready documentation and training tables
Add indexed DRL pinball documentation, provenance-aware training CSV export, first-phase flow rendering, and audited reproduction assets so the retained results are ready for manuscript analysis.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-09 22:10:05 +08:00

303 lines
16 KiB
Python

#!/usr/bin/env python3
"""Plot the validated V5 and Legacy three-role reproduction summaries."""
from __future__ import annotations
import argparse
import hashlib
import json
import os
import shutil
import tempfile
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
L0 = 20.0
VORTICITY_LIMIT = 0.001
SENSOR_PAIRS = ((0, 1, "upper"), (2, 3, "center"), (4, 5, "lower"))
SENSOR_COLORS = ("#0072B2", "#D55E00", "#009E73")
V5_GROUPS = {
"ill_075L_seed43": "ill_075L", "ill_1L_seed43": "ill_1L",
"ill_15L_seed43": "ill_15L", "ill_2L_seed43": "ill_2L",
"kar_d075_seed44": "kar_d075", "kar_d15_seed45": "kar_d15",
"kar_d2_seed45": "kar_d2", "kar_re60_seed43": "kar_re60",
"kar_re100_seed41": "kar_re100", "kar_re100_seed42": "kar_re100",
"kar_re100_seed43": "kar_re100", "kar_re100_seed44": "kar_re100",
"kar_re100_seed45": "kar_re100", "kar_re200_seed43": "kar_re200",
"kar_re400_seed43": "kar_re400",
}
LEGACY_PERIODIC = (
"illusion_075L", "illusion_1L", "illusion_15L", "karman_re50",
"karman_re100", "karman_re200", "karman_re400",
)
LEGACY_VORTEX = (
"vortex_lamb_y000", "vortex_taylor_ym2L", "vortex_taylor_ym1L",
"vortex_taylor_y000", "vortex_taylor_yp1L", "vortex_taylor_yp2L",
)
LEGACY_GROUPS = LEGACY_PERIODIC + ("steady",) + LEGACY_VORTEX + ("erase",)
SR_TRAINING = LEGACY_PERIODIC
SR_GENERALIZATION = (
"karman_re25", "karman_re70", "karman_re150", "karman_re300",
"illusion_05L", "illusion_06L", "illusion_08L", "illusion_12L", "illusion_2L",
)
SR_VARIANT_CASE = {
"k_front0": "karman_re100", "k_rear0": "karman_re100", "k_rear1": "karman_re100",
"i_front0": "illusion_15L", "i_front1": "illusion_15L", "i_rear0": "illusion_15L", "i_rear1": "illusion_15L",
}
def fail(message: str) -> None:
raise RuntimeError(message)
def require_exact_groups(root: Path) -> None:
v5_controlled = {p.parent.name for p in (root / "v5").glob("*/controlled") if p.is_dir()}
if v5_controlled != set(V5_GROUPS):
fail(f"V5 controlled groups differ: missing={set(V5_GROUPS)-v5_controlled}, extra={v5_controlled-set(V5_GROUPS)}")
legacy = {p.name for p in (root / "legacy").iterdir() if p.is_dir()}
if legacy != set(LEGACY_GROUPS):
fail(f"Legacy groups differ: missing={set(LEGACY_GROUPS)-legacy}, extra={legacy-set(LEGACY_GROUPS)}")
def roles_for(root: Path, pipeline: str, group: str):
if pipeline == "v5":
case = V5_GROUPS[group]
return (("target", root / "v5" / case / "target"),
("controlled", root / "v5" / group / "controlled"),
("physical-zero", root / "v5" / case / "zero"))
middle = "constant" if group == "steady" else "controlled"
return (("target", root / "legacy" / group / "target"),
(middle, root / "legacy" / group / middle),
("physical-zero", root / "legacy" / group / "zero"))
def field_spec(pipeline: str, group: str, role_dir: Path):
if pipeline == "v5" or group in LEGACY_PERIODIC + SR_GENERALIZATION or (group == "erase" and role_dir.name != "target"):
return "phase_fields.npz", "target_phase", 0.0
if group in LEGACY_VORTEX:
return "event_fields.npz", "relative_offsets", -10
return "late_field.npz", "field_indices", None
def load_role(role_dir: Path, pipeline: str, group: str):
if not role_dir.is_dir():
fail(f"missing role directory: {role_dir}")
required = {"metadata.json", "timeseries.csv"}
names = {p.name for p in role_dir.iterdir() if p.is_file()}
if not required <= names:
fail(f"missing required files in {role_dir}: {required-names}")
field_name, selector_name, selector_value = field_spec(pipeline, group, role_dir)
field_candidates = names & {"phase_fields.npz", "event_fields.npz", "late_field.npz"}
if field_candidates != {field_name}:
fail(f"field files differ in {role_dir}: expected {field_name}, found {sorted(field_candidates)}")
metadata = json.loads((role_dir / "metadata.json").read_text())
units = metadata.get("units", {}).get("sensors", "raw sensor units")
data = np.genfromtxt(role_dir / "timeseries.csv", delimiter=",", names=True)
if data.ndim != 1 or data.size == 0:
fail(f"invalid timeseries shape in {role_dir}: {data.shape}")
expected_sensors = {f"sensors_{i}" for i in range(6)}
actual_sensors = {n for n in (data.dtype.names or ()) if n.startswith("sensors_")}
if actual_sensors != expected_sensors:
fail(f"sensor columns differ in {role_dir}: {actual_sensors}")
sensors = np.column_stack([data[f"sensors_{i}"] for i in range(6)])
if sensors.shape != (data.size, 6) or not np.isfinite(sensors).all():
fail(f"invalid sensor values in {role_dir}: {sensors.shape}")
field_path = role_dir / field_name
with np.load(field_path) as z:
required_keys = {"ux", "uy", selector_name}
if not required_keys <= set(z.files):
fail(f"missing field keys in {field_path}: {required_keys-set(z.files)}")
ux, uy = np.array(z["ux"]), np.array(z["uy"])
selector = np.array(z[selector_name])
if ux.ndim != 3 or ux.shape != uy.shape or ux.shape[0] != selector.shape[0]:
fail(f"invalid canonical field shapes in {field_path}: ux={ux.shape}, uy={uy.shape}, selector={selector.shape}")
expected_count = 8 if field_name == "phase_fields.npz" else 5 if field_name == "event_fields.npz" else 1
if ux.shape[0] != expected_count or selector.shape != (expected_count,):
fail(f"unexpected retained field count in {field_path}: {ux.shape[0]}")
if selector_value is not None and not np.isclose(selector[0], selector_value, atol=1e-12, rtol=0):
fail(f"slot 0 selector in {field_path} is {selector[0]}, expected {selector_value}")
if not np.isfinite(ux[0]).all() or not np.isfinite(uy[0]).all():
fail(f"nonfinite field in {field_path} slot 0")
omega = np.gradient(uy[0], axis=1) - np.gradient(ux[0], axis=0)
return omega, sensors, units, field_path, selector_name, selector[0]
def padded_limits(values):
lo, hi = float(np.min(values)), float(np.max(values))
if not np.isfinite([lo, hi]).all():
fail("nonfinite sensor limits")
span = hi - lo
pad = 0.05 * span if span > 0 else max(abs(lo) * 0.05, 1e-12)
return [lo - pad, hi + pad]
def sha256_file(path: Path):
digest = hashlib.sha256()
with path.open("rb") as stream:
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def plot_role_set(stage: Path, data_root: Path, group: str, roles, output_name: str):
loaded = [load_role(path, "legacy", group) for _, path in roles]
shapes = {item[0].shape for item in loaded}
if len(shapes) != 1:
fail(f"role field shapes differ for SR/{group}: {shapes}")
all_sensors = np.concatenate([item[1] for item in loaded], axis=0)
xlim = padded_limits(all_sensors[:, [0, 2, 4]])
ylim = padded_limits(all_sensors[:, [1, 3, 5]])
ny, nx = next(iter(shapes)); extent = (0, (nx - 1) / L0, 0, (ny - 1) / L0)
fig, axes = plt.subplots(2, len(roles), figsize=(6 * len(roles), 8), constrained_layout=True, squeeze=False)
images = []
for col, ((label, _), (omega, sensors, units, _, _, _)) in enumerate(zip(roles, loaded)):
images.append(axes[0, col].imshow(omega, origin="lower", extent=extent, cmap="RdBu_r", vmin=-VORTICITY_LIMIT, vmax=VORTICITY_LIMIT, aspect="equal"))
axes[0, col].set_title(f"LEGACY / {group}{label}")
axes[0, col].set_xlabel("x/L0"); axes[0, col].set_ylabel("y/L0")
axes[0, col].set_xlim(extent[:2]); axes[0, col].set_ylim(extent[2:])
for (u, v, sensor_label), color in zip(SENSOR_PAIRS, SENSOR_COLORS):
axes[1, col].plot(sensors[:, u], sensors[:, v], color=color, lw=1.1, alpha=0.85, label=sensor_label)
axes[1, col].axhline(0, color="0.35", lw=0.7); axes[1, col].axvline(0, color="0.35", lw=0.7)
axes[1, col].grid(True, alpha=0.25); axes[1, col].set_xlim(xlim); axes[1, col].set_ylim(ylim)
axes[1, col].set_xlabel(f"sensor u ({units})"); axes[1, col].set_ylabel(f"sensor v ({units})")
axes[1, col].set_box_aspect(0.8); axes[1, col].legend(loc="best", frameon=False)
cbar = fig.colorbar(images[0], ax=axes[0, :], orientation="vertical", shrink=0.92, pad=0.015)
cbar.set_label(r"$\omega_z$ (lattice$^{-1}$)")
output = stage / "legacy" / f"{output_name}.png"; output.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(output, dpi=160); plt.close(fig)
role_entries = []
for (label, path), item in zip(roles, loaded):
sources = [path / "metadata.json", path / "timeseries.csv", item[3]]
role_entries.append({"role": label, "source": str(path), "field_file": str(item[3]), "field_slot": 0,
"selector": item[4], "selector_value": float(item[5]),
"source_files": [{"path": str(source), "sha256": sha256_file(source), "bytes": source.stat().st_size} for source in sources]})
return {"pipeline": "legacy-sr", "group": group, "plot": f"legacy/{output_name}.png", "roles": role_entries,
"vorticity_limit": VORTICITY_LIMIT, "sensor_limits": {"u": xlim, "v": ylim}}
def sr_plot_specs(data_root: Path):
legacy = data_root / "legacy"; specs = []
for case in SR_TRAINING:
base = legacy / case
specs.append((case, (("Target", base / "target"), ("PPO", base / "controlled"), ("SR", base / "sr"), ("Zero (physical, no control)", base / "zero")), case))
for case in SR_GENERALIZATION:
base = legacy / case
specs.append((case, (("Target", base / "target"), ("SR", base / "sr"), ("Zero (physical, no control)", base / "zero")), case))
for variant, case in SR_VARIANT_CASE.items():
base = legacy / case
specs.append((case, (("Target", base / "target"), ("Parent SR", base / "sr"), (f"Variant SR ({variant})", base / f"sr_{variant}"), ("Zero (physical, no control)", base / "zero")), f"variants/{case}_{variant}"))
return specs
def render_sr_package(stage: Path, data_root: Path, allow_partial=False):
entries = []
for group, roles, output_name in sr_plot_specs(data_root):
missing = [str(path) for _, path in roles if not path.is_dir()]
if missing:
if allow_partial:
continue
fail(f"missing SR plot roles for {group}: {missing}")
entries.append(plot_role_set(stage, data_root, group, roles, output_name))
if not allow_partial and len(entries) != 23:
fail(f"expected 23 SR plots, got {len(entries)}")
return entries
def plot_group(stage: Path, data_root: Path, pipeline: str, group: str):
roles = roles_for(data_root, pipeline, group)
loaded = [load_role(path, pipeline, group) for _, path in roles]
shapes = {x[0].shape for x in loaded}
if len(shapes) != 1:
fail(f"role field shapes differ for {pipeline}/{group}: {shapes}")
limit = VORTICITY_LIMIT
all_sensors = np.concatenate([x[1] for x in loaded], axis=0)
xlim = padded_limits(all_sensors[:, [0, 2, 4]])
ylim = padded_limits(all_sensors[:, [1, 3, 5]])
ny, nx = next(iter(shapes))
extent = (0, (nx - 1) / L0, 0, (ny - 1) / L0)
fig, axes = plt.subplots(2, 3, figsize=(18, 8), constrained_layout=True)
images = []
for col, ((role, _), (omega, sensors, units, _, _, _)) in enumerate(zip(roles, loaded)):
images.append(axes[0, col].imshow(omega, origin="lower", extent=extent, cmap="RdBu_r", vmin=-limit, vmax=limit, aspect="equal"))
axes[0, col].set_title(f"{pipeline.upper()} / {group}{role}")
axes[0, col].set_xlabel("x/L0")
axes[0, col].set_ylabel("y/L0")
axes[0, col].set_xlim(extent[:2]); axes[0, col].set_ylim(extent[2:])
for (u, v, label), color in zip(SENSOR_PAIRS, SENSOR_COLORS):
axes[1, col].plot(sensors[:, u], sensors[:, v], color=color, lw=1.1, alpha=0.85, label=label)
axes[1, col].axhline(0, color="0.35", lw=0.7); axes[1, col].axvline(0, color="0.35", lw=0.7)
axes[1, col].grid(True, alpha=0.25)
axes[1, col].set_xlim(xlim); axes[1, col].set_ylim(ylim)
axes[1, col].set_xlabel(f"sensor u ({units})"); axes[1, col].set_ylabel(f"sensor v ({units})")
axes[1, col].set_box_aspect(0.8)
axes[1, col].legend(loc="best", frameon=False)
cbar = fig.colorbar(images[0], ax=axes[0, :], orientation="vertical", shrink=0.92, pad=0.015)
cbar.set_label(r"$\omega_z$ (lattice$^{-1}$)")
output = stage / pipeline / f"{group}.png"
output.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(output, dpi=160)
plt.close(fig)
return {
"pipeline": pipeline, "group": group, "plot": f"{pipeline}/{group}.png",
"roles": [{"role": role, "source": str(path), "field_file": str(item[3]),
"field_slot": 0, "selector": item[4], "selector_value": float(item[5])}
for (role, path), item in zip(roles, loaded)],
"vorticity_limit": limit, "sensor_limits": {"u": xlim, "v": ylim},
}
def main() -> None:
here = Path(__file__).resolve().parent
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--data-root", type=Path, default=here / "data" / "reproduction")
parser.add_argument("--output-root", type=Path)
parser.add_argument("--campaign", choices=("ppo", "sr"), default="ppo")
parser.add_argument("--allow-partial", action="store_true")
args = parser.parse_args()
data_root = args.data_root.resolve()
output_root = (args.output_root or here / "data" / ("reproduction_plots_sr" if args.campaign == "sr" else "reproduction_plots")).absolute()
if args.campaign == "ppo": require_exact_groups(data_root)
if output_root.is_symlink():
fail(f"refusing symlink output root: {output_root}")
output_root.parent.mkdir(parents=True, exist_ok=True)
stage = Path(tempfile.mkdtemp(prefix=f".{output_root.name}.staging-", dir=output_root.parent))
backup = output_root.with_name(f".{output_root.name}.old-{os.getpid()}")
try:
if args.campaign == "sr":
entries = render_sr_package(stage, data_root, allow_partial=args.allow_partial)
manifest = {"schema": "drl-pinball-sr-reproduction-plots-v1", "plot_count": len(entries),
"expected_plot_count": 23, "allow_partial": args.allow_partial,
"no_ppo_cases": list(SR_GENERALIZATION),
"no_ppo_contract": "Generalization conditions intentionally contain Target/SR/Zero only; PPO is absent by acquisition design.",
"variant_contract": "Each diagnostic binds the named variant to its same-case parent SR.",
"vorticity_limit": VORTICITY_LIMIT,
"vorticity_contract": "fixed symmetric [-0.001, +0.001] for every vorticity panel",
"plots": entries}
else:
entries = [plot_group(stage, data_root, "v5", group) for group in sorted(V5_GROUPS)]
entries += [plot_group(stage, data_root, "legacy", group) for group in LEGACY_GROUPS]
if len(entries) != 30:
fail(f"expected 30 plots, got {len(entries)}")
manifest = {"schema": "drl-pinball-reproduction-plots-v1", "plot_count": 30,
"vorticity_limit": VORTICITY_LIMIT,
"vorticity_contract": "fixed symmetric [-0.001, +0.001] for every vorticity panel", "plots": entries}
(stage / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
if output_root.exists():
os.replace(output_root, backup)
os.replace(stage, output_root)
if backup.exists():
shutil.rmtree(backup)
except Exception:
shutil.rmtree(stage, ignore_errors=True)
if backup.exists() and not output_root.exists():
os.replace(backup, output_root)
raise
print(f"wrote {len(entries)} plots and manifest to {output_root}")
if __name__ == "__main__":
main()