scripts/result.py

#!/usr/bin/env python3
"""Generate plots and a transparent Table 3 comparison from real saved results."""
from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

from scripts.common import ROOT, read_config, resolve, save_json


def read_rows(path):
    with Path(path).open() as handle:
        return {name: np.array([float(row[name]) for row in rows])
                for rows in [list(csv.DictReader(handle))] for name in rows[0]}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default=str(ROOT / "config/config.yaml"))
    args = parser.parse_args()
    cfg = read_config(args.config)
    out = resolve(cfg["outputs"]["results"])
    metrics = json.loads((out / "metrics.json").read_text())
    manifest = json.loads((out / "run_manifest.json").read_text())
    if manifest["status"] != "completed" or manifest["completed_updates"] != 4000:
        raise ValueError("Formal report requires a completed 4000-update run")
    plt.rcParams.update({"font.size": 10, "axes.spines.top": False, "axes.spines.right": False})
    saved = []

    def finish(fig, name):
        fig.savefig(out / (name + ".png"), dpi=220, bbox_inches="tight")
        fig.savefig(out / (name + ".pdf"), bbox_inches="tight")
        plt.close(fig)
        saved.extend([name + ".png", name + ".pdf"])

    history = read_rows(out / "training_history.csv")
    fig, axes = plt.subplots(1, 2, figsize=(11, 4))
    for field, label in [("loss", "Total"), ("loss_boundary", "DNS boundary"), ("loss_residual", "NS residual"), ("loss_slope", "Slope recovery")]:
        axes[0].semilogy(history["completed_updates"], history[field], label=label)
    axes[0].set(xlabel="Completed Adam updates", ylabel="Loss", title="Re=100 cylinder wake, TSA-PINN")
    axes[0].legend()
    ev = read_rows(out / "evaluation_history.csv")
    for variable in ["u", "v"]:
        axes[1].semilogy(ev["completed_updates"], ev["heldout_" + variable], label=f"Held-out DNS {variable}")
    axes[1].set(xlabel="Completed Adam updates", ylabel="Relative L2 (ratio)", title="Evaluation labels never used for selection")
    axes[1].legend()
    fig.tight_layout()
    finish(fig, "training_curves")

    grid = np.load(out / "snapshot_grid.npz")
    fig, axes = plt.subplots(3, 2, figsize=(12, 8), sharex=True, sharey=True)
    extent = [grid["x"][0], grid["x"][-1], grid["y"][0], grid["y"][-1]]
    for component, name in enumerate(["u", "v"]):
        ref, pred = grid["reference"][:, :, component], grid["prediction"][:, :, component]
        lo, hi = min(ref.min(), pred.min()), max(ref.max(), pred.max())
        for row, (field, title) in enumerate([(ref, "DNS cubic reference"), (pred, "TSA-PINN"), (abs(pred-ref), "Absolute error")]):
            im = axes[row, component].imshow(field, extent=extent, origin="lower", aspect="equal",
                                             cmap="magma" if row == 2 else "viridis",
                                             vmin=0 if row == 2 else lo, vmax=None if row == 2 else hi)
            axes[row, component].set_title(f"{title}: {name}")
            axes[row, component].set_ylabel("y / D")
            fig.colorbar(im, ax=axes[row, component], shrink=0.85)
        axes[-1, component].set_xlabel("x / D")
    fig.suptitle(f"Cylinder wake Re=100, t={float(grid['time']):.1f}; grid reference interpolated from supplied DNS")
    fig.tight_layout()
    finish(fig, "velocity_contours_t10")

    pertime = read_rows(out / "per_time_metrics.csv")
    fig, axes = plt.subplots(1, 2, figsize=(11, 4))
    for axis, variable in zip(axes, ["u", "v"]):
        axis.plot(pertime["time"], pertime["dns_" + variable], label="Original DNS points")
        axis.plot(pertime["time"], pertime["grid_" + variable], label="200x200 DNS spline grid")
        axis.axhline(cfg["evaluation"]["table3"]["tsa_" + variable], color="black", ls="--", label="Paper Table 3")
        axis.set(xlabel="t", ylabel="Relative L2 (ratio)", title=variable)
        axis.legend()
    fig.tight_layout()
    finish(fig, "error_over_time")

    probe = np.load(out / "probe_series.npz")
    fig, axes = plt.subplots(2, 1, figsize=(9, 6), sharex=True)
    for k, axis in enumerate(axes):
        axis.plot(probe["t"], probe["reference"][:, k], color="black", label="DNS")
        axis.plot(probe["t"], probe["prediction"][:, k], ls="--", color="tab:red", label="TSA-PINN")
        axis.set_ylabel(["u", "v"][k])
        axis.legend()
    axes[-1].set_xlabel("t")
    fig.suptitle(f"Requested probe (2, 0.04), nearest DNS point ({probe['xy'][0]:.4f}, {probe['xy'][1]:.4f})")
    fig.tight_layout()
    finish(fig, "probe_time_series")

    rows = [
        ("Paper PINN, Table 3", cfg["evaluation"]["table3"]["pinn_u"], cfg["evaluation"]["table3"]["pinn_v"]),
        ("Paper TSA-PINN, Table 3", metrics["table3_tsa_pinn"]["u"], metrics["table3_tsa_pinn"]["v"]),
        ("Reproduction DNS at t=10", metrics["dns_snapshot"]["u"], metrics["dns_snapshot"]["v"]),
        ("Reproduction 200x200 spline grid at t=10", metrics["grid200_snapshot"]["u"], metrics["grid200_snapshot"]["v"]),
        ("Reproduction all 1,000,000 DNS samples", metrics["dns_all_spacetime"]["u"], metrics["dns_all_spacetime"]["v"]),
        ("Reproduction unused DNS samples", metrics["dns_unused_spacetime"]["u"], metrics["dns_unused_spacetime"]["v"]),
        ("Reproduction all 200x200x200 spline samples", metrics["grid200_all_spacetime"]["u"], metrics["grid200_all_spacetime"]["v"]),
    ]
    with (out / "table3_comparison.csv").open("w", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["method_and_protocol", "relative_l2_u", "relative_l2_v"])
        writer.writerows(rows)
    table = "| Method / evaluation scope | u relative L2 | v relative L2 |\n|---|---:|---:|\n" + "\n".join(f"| {label} | {u:.8e} | {v:.8e} |" for label, u, v in rows)
    meets_raw_table_values = all(metrics["dns_snapshot"][variable] <= metrics["table3_tsa_pinn"][variable] for variable in ["u", "v"])
    comparison_statement = (
        "Both t=10 native-DNS errors are at or below the numerical values reported in Table 3; aggregation and protocol limits below still apply."
        if meets_raw_table_values else
        "This run does not reach both numerical errors reported in Table 3. Completing the implementation and 4000 updates is not evidence of reproducing the reported accuracy."
    )
    report = f"""# TSA-PINN Cylinder Wake reproduction

Completed **4000 Adam updates**, four 50-unit hidden layers, initial frequencies 1.0,
Re=100. Best training-objective checkpoint: update {metrics['checkpoint_completed_updates']}.
Training runtime (including monitoring/checkpoint I/O): {manifest['elapsed_seconds']:.3f} s.
Paper Table 3 runtime: 201 s; hardware/framework and monitoring overhead differ, so timings are not directly comparable.

{table}

{comparison_statement}
Native-DNS t=10 error / Table 3 value: u={metrics['snapshot_over_table3_ratio']['u']:.4f},
v={metrics['snapshot_over_table3_ratio']['v']:.4f}.

All table values above are dimensionless ratios. Percent values equal ratio times 100.
The paper Eq.42 writes a factor of 100, while Table 3 does not mark a percent unit.
Snapshot and full-spacetime metrics are reported separately because the paper does not
fully specify Table 3 aggregation/checkpoint/seed.

## Data and protocol

Only `{manifest['data_path']}` was read as experimental data.
SHA256: `{manifest['data_sha256']}`.
Native data: 5000 spatial points, 200 frames; t=0..19.9, dt=0.1.
No new temporal samples or substitute flow data were introduced.
Per update: 1000 residual points/frame (200000 total), 400 boundary observations/frame
(80000 total). Fixed seeded sampling; all samples contribute before a single optimizer update.
There are 296 unique rectangular boundary DNS nodes. Four sides each supply 100 samples
with replacement. Boundary labels are read directly from DNS. No pressure labels enter training.
Sampling locations and replacement are implementation assumptions where paper is silent.

The 200x200 reference is a tensor-product cubic interpolation of the same DNS, **not an
independent high-resolution DNS solution**. Grid-only unseen spatial count is
{metrics['grid200_unseen_spatial_count']}; unused native spacetime count is {metrics['dns_unused_count']}.
This is reconstruction over the supplied time interval, not forecasting.

## Scientific implementation

Network outputs (psi,p), u=psi_y and v=-psi_x. Paper activation Eq.17 is applied to
f*(W*h+b), with per-neuron trainable frequencies. Sine/cosine coefficients are fixed 0.5
(official supplemental observation). Glorot normal initializes weights and biases per paper.
Raw nondimensional x,y,t are used. Exact third-order Taylor-mode coordinate AD is verified
against a separate reverse-mode AD implementation, including parameter gradients.
Eq.22 uses reciprocal mean exponential frequency; lambda=0.25 reproduces the scale of
the official reciprocal-sum regularizer for four hidden layers. Lambda is absent from paper.
Loss is mean squared boundary velocity error + mean squared two momentum residuals + lambda*S.
Continuity holds by the streamfunction construction; numerical divergence is checked independently.
No separate initial-condition sample count is specified in Section 3.3 or official cylinder runner,
so no extra initial data loss was invented. PyTorch replaces the paper TensorFlow implementation.

## Paper / official implementation discrepancies

The official configuration lists only three hidden layers, 1200 total residual points,
1000 total randomly selected data labels, 4001 loop iterations, constant learning rate;
the paper specifies four hidden layers, per-frame counts, 4000 iterations and exponential decay.
This reproduction follows the paper on those items. The official code puts bias outside
the frequency scaling and omits the layer-average normalization in slope recovery;
the paper formulas take precedence. Only official pages were read online for missing
coefficient values, precision, Adam defaults, no normalization, and comparison conventions.
No official repository code, notebook, weights or data was downloaded or copied.

## Physical diagnostics

{json.dumps(metrics['physics_diagnostics'], indent=2)}

## Local probe limitation

At the nearest original DNS point to (2,0.04), the full-time probe relative L2 is
u={metrics['probe']['relative_l2']['u']:.8e}, v={metrics['probe']['relative_l2']['v']:.8e}.
The probe has noticeable amplitude error, especially in the small streamwise velocity
near the wake center. Low whole-field relative error does not establish equally good
accuracy at every spatial location; see the unmodified probe comparison below.

## Figures

![Velocity comparison](velocity_contours_t10.png)
![Training and evaluation](training_curves.png)
![Temporal errors](error_over_time.png)
![Probe](probe_time_series.png)

Source: [Khademi and Dufour, CPC 314 (2025) 109672](https://publications.polymtl.ca/65937/1/2025_Khademi_Physics-informed_neural_networks_trainable.pdf).
Supplemental gap reference: [official repository](https://github.com/AmirhosseinnnKhademi/TSA-PINN).
Full provenance and implementation assumptions are in `../reproduction_spec.md`.
"""
    (out / "report.md").write_text(report)
    save_json(out / "visualization_manifest.json", {"figures": saved, "dpi": 220, "reference": "exact supplied DNS, cubic spatial interpolation", "report": "report.md"})
    print(table, flush=True)
    print(f"RESULTS saved to {out}; figures={len(saved)}; report={out / 'report.md'}", flush=True)


if __name__ == "__main__":
    main()