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>
303 lines
16 KiB
Python
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()
|