| """Create standalone scientific figures and audit actual inference artifacts.""" |
| import json |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
|
|
| from common import FIELDS, ROOT, load_config, output_dir, parser, save_json, sha256 |
|
|
|
|
| def main(): |
| args = parser(__doc__).parse_args() |
| cfg = load_config(args.config) |
| out = output_dir(cfg) |
| history = json.loads((out / "training_history.json").read_text()) |
| status = json.loads((out / "training_status.json").read_text()) |
| manifest = json.loads((out / "inference_manifest.json").read_text()) |
| metrics = json.loads((out / "metrics.json").read_text()) |
| assert status["completed_epochs"] == 300 and status["state"] == "completed" |
| assert [r["epoch"] for r in history] == list(range(1, 301)) |
| assert all(np.isfinite(r["train_loss"]) and np.isfinite(r["val_mse"]) for r in history) |
| selected = min(history, key=lambda r: r["val_mse"]) |
| assert selected["epoch"] == manifest["selected_epoch"] == status["best_epoch"] |
| assert sha256(ROOT / cfg["paths"]["best_checkpoint"]) == manifest["sha256"] |
| figures = [] |
| plt.rcParams.update({"font.size": 10, "figure.dpi": 140, "savefig.dpi": 180}) |
| fig, ax = plt.subplots(figsize=(7, 4)) |
| for key, label in [("train_step1_mse", "Train step 1"), |
| ("train_step5_mse", "Train step 5"), ("val_mse", "Validation step 1 + 5")]: |
| ax.semilogy([r["epoch"] for r in history], [r[key] for r in history], label=label) |
| ax.axvline(status["best_epoch"], color="0.6", lw=1, ls=":", label="Selected checkpoint") |
| ax.set(xlabel="Epoch", ylabel="MSE in physical units", title="FluxNet-LAP: 300-epoch pushforward training") |
| ax.legend() |
| fig.tight_layout() |
| fig.savefig(out / "training_curve.png") |
| plt.close(fig) |
| figures.append("training_curve.png") |
| checks = {} |
| for split in ("test", "test_long"): |
| data = np.load(out / f"{split}.npz") |
| expected_steps = cfg["data"]["expected_frames"][split] - 1 |
| assert data["mae"].shape == (50, expected_steps, 3) |
| assert np.isfinite(data["mae"]).all() |
| assert np.allclose(data["mae"].mean((0, 1)), metrics[split]["mae"], atol=1e-12) |
| assert np.allclose(data["mae"][:, -1].mean(0), metrics[split]["mae_final"], atol=1e-12) |
| assert float(data["negative_h_ratio"].mean()) == metrics[split]["negative_h_ratio"] |
| assert np.allclose(data["conservation_absolute_sum"].max((0, 1)), |
| metrics[split]["conservation_absolute_sum_max"]) |
| fig, axes = plt.subplots(1, 3, figsize=(13, 3.6)) |
| for c, field in enumerate(FIELDS): |
| axes[0].plot(data["time"], data["mae"].mean(0)[:, c], label=field) |
| axes[1].plot(data["time"], data["conservation_absolute_sum"].max(0)[:, c], label=field) |
| axes[0].set(xlabel="Time", ylabel="MAE", title="Autoregressive error") |
| axes[1].set(xlabel="Time", ylabel="max |sum(u)-sum(u0)|", title="Conservation drift (float64 reduction)") |
| axes[2].plot(data["time"], 100 * data["negative_h_ratio"].mean(0)) |
| axes[2].set(xlabel="Time", ylabel="Negative depth (%)", title="Depth bound violations") |
| for ax in axes[:2]: |
| ax.legend() |
| fig.tight_layout() |
| name = f"{split}_rollout_metrics.png" |
| fig.savefig(out / name) |
| figures.append(name) |
| plt.close(fig) |
| examples = np.load(out / f"{split}_examples.npz") |
| for case in cfg["evaluation"]["example_cases"]: |
| prediction = examples[case + "_prediction"] |
| target = examples[case + "_target"] |
| assert prediction.shape == target.shape == (expected_steps + 1, 3, 64, 64) |
| assert np.array_equal(prediction[0], target[0]) |
| |
| errors = np.abs(prediction[1:].astype(np.float64) - target[1:]).mean((0, 2, 3)) |
| file = str(examples[case + "_file"]) |
| index = list(data["files"]).index(file) |
| assert np.allclose(errors, data["mae"][index].mean(0), atol=1e-12) |
| frames = np.round(np.linspace(0, expected_steps, 6)).astype(int) |
| for c, field in enumerate(FIELDS): |
| fig, axes = plt.subplots(3, 6, figsize=(14, 7.3), layout="constrained") |
| minimum = min(target[frames, c].min(), prediction[frames, c].min()) |
| maximum = max(target[frames, c].max(), prediction[frames, c].max()) |
| error_max = np.abs(prediction[frames, c] - target[frames, c]).max() |
| for j, frame in enumerate(frames): |
| for r, values in enumerate((target, prediction)): |
| image = axes[r, j].imshow(values[frame, c], origin="lower", extent=(0, 10, 0, 10), |
| cmap="viridis", vmin=minimum, vmax=maximum) |
| err_image = axes[2, j].imshow(np.abs(prediction[frame, c] - target[frame, c]), |
| origin="lower", extent=(0, 10, 0, 10), cmap="magma", |
| vmin=0, vmax=max(float(error_max), 1e-12)) |
| axes[0, j].set_title(f"t={examples['time'][frame]:.2f}") |
| for r in range(3): |
| axes[r, j].set_xticks([]) |
| axes[r, j].set_yticks([]) |
| for r, label in enumerate(("Reference", "FluxNet-LAP", "Absolute error")): |
| axes[r, 0].set_ylabel(label) |
| fig.colorbar(image, ax=axes[:2, :], shrink=0.75) |
| fig.colorbar(err_image, ax=axes[2, :], shrink=0.75) |
| fig.suptitle(f"{split} Case {case}: {field}") |
| name = f"{split}_{case}_{field}.png" |
| fig.savefig(out / name) |
| figures.append(name) |
| plt.close(fig) |
| checks[split] = {"all_trajectories": True, "all_steps": True, "finite": True, |
| "metrics_recomputed": True, "example_fields_recomputed": True, |
| "zero_negative_depth": metrics[split]["negative_h_ratio"] == 0} |
| report = ["# FluxNet shallow-water LAP pushforward reproduction", "", |
| "Source: https://arxiv.org/pdf/2602.01941v2 (Sections 3.2–3.7, Appendix C.2).", "", |
| "One independent training run (seed 42), 300 complete epochs. This is not a five-seed mean/std replication.", |
| "64 channels, 6 residual blocks, kernel 5, transport stencil 3×3, detached pushforward length 5.", |
| "AdamW lr=1e-3, weight_decay=1e-2; ReduceLROnPlateau factor=0.5, patience=15; float32 states.", "", |
| f"Selected checkpoint epoch: {manifest['selected_epoch']}; validation MSE: {manifest['best_val_mse']:.8e}.", |
| f"Checkpoint SHA256: `{manifest['sha256']}`.", "", |
| "| Split | Horizon | MAE h | MAE mx | MAE my | Negative h (%) |", |
| "| --- | ---: | ---: | ---: | ---: | ---: |"] |
| for split, values in metrics.items(): |
| report.append(f"| {split} | {values['t_end']:.1f} | " + " | ".join(f"{v:.8e}" for v in values["mae"]) + |
| f" | {values['negative_h_percent']:.8e} |") |
| report += ["", "MAE averages all forecast steps, grid cells, and 50 trajectories, excluding the initial frame.", |
| "Final-time and T=2/T=2.4/T=4.8 prefix metrics are separately stored in metrics.json.", |
| "The paper labels Table 3 as T=2 while Appendix C.1 specifies 2.4/4.8; both time conventions are retained.", |
| "", "Conservation is computed using float64 reductions of float32 predictions relative to the initial state.", |
| "Relative drift divides by |initial signed total|+1e-8; zero-total momentum makes this statistic ill-conditioned.", |
| "Absolute sum drift and initial-L1-scaled drift are therefore also reported. No postprocessing correction is used.", ""] |
| for split, values in metrics.items(): |
| report += [f"{split} maximum absolute sum drift [h,mx,my]: {values['conservation_absolute_sum_max']}", |
| f"{split} maximum relative drift [h,mx,my]: {values['conservation_relative_max']}", |
| f"{split} minimum depth: {values['min_h']:.8e}", ""] |
| report += ["All arrays, checkpoints, training history, and saved-case MAE recomputations were checked.", |
| "The original numerical-generator datasets were read directly; no data generation or official repository code was downloaded."] |
| report += ["", "Reference data contains small negative roundoff depths (about 1e-11); these are retained unchanged.", |
| "Data validation tolerates reference values down to -1e-9; model violation metrics use the strict condition h<0, with no tolerance or clipping."] |
| (out / "report.md").write_text("\n".join(report) + "\n") |
| save_json(out / "verification.json", {"completed_epochs": 300, |
| "best_checkpoint_hash_matches": True, "selected_epoch_is_val_minimum": True, |
| "checkpoint_sha256": manifest["sha256"], "splits": checks, |
| "figures": figures, "report": "report.md"}) |
| print("RESULT_VERIFIED " + json.dumps({"completed_epochs": 300, "checks": checks, |
| "figures": len(figures), "results": str(out)}), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
