| |
| """Reload best checkpoint; evaluate exact DNS and DNS-interpolated 200x200 grids.""" |
| 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 numpy as np |
| import torch |
| from scipy.interpolate import RectBivariateSpline |
|
|
| from scripts.common import (ROOT, gather, load_checkpoint, load_dns, predict, read_config, |
| relative_l2, resolve, save_json, setup, tensor) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--config", default=str(ROOT / "config/config.yaml")) |
| parser.add_argument("--checkpoint") |
| parser.add_argument("--device") |
| args = parser.parse_args() |
| cfg = read_config(args.config) |
| device, dtype = setup(cfg, args.device) |
| model, checkpoint = load_checkpoint(cfg, device, args.checkpoint) |
| dns = load_dns(cfg) |
| if checkpoint["data_sha256"] != dns["sha256"]: |
| raise ValueError("Checkpoint was trained on different DNS bytes") |
| out = resolve(cfg["outputs"]["results"]) |
| sampling = np.load(out / "sampling_indices.npz") |
| xy, times = dns["xy"], dns["t"] |
| nx, ny = cfg["evaluation"]["grid_size"] |
| gx = np.linspace(xy[:, 0].min(), xy[:, 0].max(), nx) |
| gy = np.linspace(xy[:, 1].min(), xy[:, 1].max(), ny) |
| xx, yy = np.meshgrid(gx, gy) |
| grid_xy = np.column_stack((xx.ravel(), yy.ravel())) |
| native_x, native_y = np.unique(xy[:, 0]), np.unique(xy[:, 1]) |
| ix, iy = np.searchsorted(native_x, xy[:, 0]), np.searchsorted(native_y, xy[:, 1]) |
| dns_grid_coincidence = np.any(np.isclose(gx[:, None], native_x[None, :], atol=1e-10, rtol=0), 1)[None, :] & np.any(np.isclose(gy[:, None], native_y[None, :], atol=1e-10, rtol=0), 1)[:, None] |
| grid_unseen = ~dns_grid_coincidence.ravel() |
| native_pred = np.empty((len(times), len(xy), 3), dtype=np.float64) |
| per_time, grid_sse, grid_denom = [], np.zeros(2), np.zeros(2) |
| unseen_sse, unseen_denom = np.zeros(2), np.zeros(2) |
| held_sse, held_denom = np.zeros(2), np.zeros(2) |
| snapshot = int(np.argmin(abs(times - cfg["evaluation"]["snapshot_time"]))) |
| batch_size = cfg["evaluation"]["batch_size"] |
| native_heldout_count = 0 |
| for k, t in enumerate(times): |
| q = np.column_stack((xy, np.full(len(xy), t))) |
| pred = predict(model, q, device, dtype, batch_size) |
| native_pred[k] = pred |
| errors = relative_l2(pred[:, :2], dns["uv"][k]) |
| unused = np.ones(len(xy), dtype=bool) |
| unused[sampling["boundary"][k]] = False |
| unused[sampling["residual"][k]] = False |
| native_heldout_count += int(unused.sum()) |
| held_sse += np.sum((pred[unused, :2] - dns["uv"][k, unused])**2, 0) |
| held_denom += np.sum(dns["uv"][k, unused]**2, 0) |
| gq = np.column_stack((grid_xy, np.full(len(grid_xy), t))) |
| gp = predict(model, gq, device, dtype, batch_size)[:, :2] |
| reference = np.empty((ny, nx, 2), dtype=np.float64) |
| for channel in range(2): |
| field = np.empty((len(native_y), len(native_x))) |
| field[iy, ix] = dns["uv"][k, :, channel] |
| reference[:, :, channel] = RectBivariateSpline(native_y, native_x, field, kx=3, ky=3)(gy, gx) |
| ref = reference.reshape(-1, 2) |
| ge = relative_l2(gp, ref) |
| grid_sse += np.sum((gp - ref)**2, 0) |
| grid_denom += np.sum(ref**2, 0) |
| unseen_sse += np.sum((gp[grid_unseen] - ref[grid_unseen])**2, 0) |
| unseen_denom += np.sum(ref[grid_unseen]**2, 0) |
| row = {"time_index": k, "time": float(t), "dns_u": errors[0], "dns_v": errors[1], "grid_u": ge[0], "grid_v": ge[1]} |
| per_time.append(row) |
| if k == snapshot: |
| np.savez_compressed(out / "snapshot_grid.npz", x=gx, y=gy, reference=reference, |
| prediction=gp.reshape(ny, nx, 2), time=t, |
| unseen_mask=grid_unseen.reshape(ny, nx)) |
| if k % 10 == 0 or k == len(times) - 1: |
| print(f"INFERENCE time_index={k:3d}/199 t={t:.2f} DNS_relL2(u,v)=({errors[0]:.8e},{errors[1]:.8e}) grid_relL2(u,v)=({ge[0]:.8e},{ge[1]:.8e})", flush=True) |
| np.savez_compressed(out / "dns_predictions.npz", X_star=xy, t=times, uvp_prediction=native_pred, |
| uv_reference=dns["uv"], p_reference=dns["p"]) |
| with (out / "per_time_metrics.csv").open("w", newline="") as handle: |
| writer = csv.DictWriter(handle, fieldnames=list(per_time[0])) |
| writer.writeheader() |
| writer.writerows(per_time) |
| primary = relative_l2(native_pred[:, :, :2].reshape(-1, 2), dns["uv"].reshape(-1, 2)) |
| snap_error = np.array([per_time[snapshot]["dns_u"], per_time[snapshot]["dns_v"]]) |
| grid_aggregate = np.sqrt(grid_sse / grid_denom) |
| grid_snapshot = np.array([per_time[snapshot]["grid_u"], per_time[snapshot]["grid_v"]]) |
| table = cfg["evaluation"]["table3"] |
| table_values = np.array([table["tsa_u"], table["tsa_v"]]) |
| nearest = int(np.argmin(np.sum((xy - np.array(cfg["evaluation"]["probe"]))**2, 1))) |
| probe_q = np.column_stack((np.tile(xy[nearest], (len(times), 1)), times)) |
| probe_prediction = predict(model, probe_q, device, dtype, batch_size) |
| np.savez_compressed(out / "probe_series.npz", xy=xy[nearest], requested_xy=cfg["evaluation"]["probe"], |
| t=times, reference=dns["uv"][:, nearest], prediction=probe_prediction[:, :2]) |
| diagnostics_q, _ = gather(dns, sampling["validation"]) |
| diagnostic_count = min(len(diagnostics_q), 1000) |
| diagnostics_q = diagnostics_q[np.linspace(0, len(diagnostics_q) - 1, diagnostic_count, dtype=int)] |
| momentum_sse, div_sse = np.zeros(2), 0.0 |
| ad_max_abs = {"flow": 0.0, "momentum": 0.0} |
| |
| for start in range(0, diagnostic_count, 250): |
| q = tensor(diagnostics_q[start:start + 250], device, dtype) |
| flow_values, residual, div = model.residual_autograd(q) |
| with torch.no_grad(): |
| taylor_values, taylor_residual, _ = model.residual_taylor(q) |
| torch.testing.assert_close(taylor_values, flow_values, atol=1e-8, rtol=1e-7) |
| torch.testing.assert_close(taylor_residual, residual, atol=1e-8, rtol=1e-7) |
| ad_max_abs["flow"] = max(ad_max_abs["flow"], float((taylor_values-flow_values).abs().max().detach())) |
| ad_max_abs["momentum"] = max(ad_max_abs["momentum"], float((taylor_residual-residual).abs().max().detach())) |
| momentum_sse += (residual.detach().cpu().numpy()**2).sum(0) |
| div_sse += float((div.detach()**2).sum()) |
| def pair(array): |
| return {"u": float(array[0]), "v": float(array[1])} |
| metrics = {"checkpoint_completed_updates": checkpoint["completed_updates"], "data_path": dns["path"], |
| "data_sha256": dns["sha256"], "metric": "||prediction-reference||_2 / ||reference||_2 (dimensionless ratio)", |
| "dns_all_spacetime": pair(primary), "dns_all_spacetime_percent": pair(primary * 100), |
| "dns_unused_spacetime": pair(np.sqrt(held_sse / held_denom)), "dns_unused_count": native_heldout_count, |
| "dns_snapshot": {"time": float(times[snapshot]), **pair(snap_error)}, |
| "grid200_all_spacetime": pair(grid_aggregate), "grid200_snapshot": {"time": float(times[snapshot]), **pair(grid_snapshot)}, |
| "grid200_strictly_unseen_spacetime": pair(np.sqrt(unseen_sse / unseen_denom)), |
| "grid200_unseen_spatial_count": int(grid_unseen.sum()), |
| "grid_reference": "Tensor-product cubic spline interpolation of this exact DNS; not additional DNS ground truth", |
| "table3_tsa_pinn": pair(table_values), "table3_pinn": {"u": table["pinn_u"], "v": table["pinn_v"]}, |
| "snapshot_over_table3_ratio": pair(snap_error / table_values), |
| "grid_snapshot_over_table3_ratio": pair(grid_snapshot / table_values), |
| "all_spacetime_over_table3_ratio": pair(primary / table_values), |
| "probe": {"requested": cfg["evaluation"]["probe"], "actual_dns_point": xy[nearest].tolist(), |
| "relative_l2": pair(relative_l2(probe_prediction[:, :2], dns["uv"][:, nearest]))}, |
| "physics_diagnostics": {"count": diagnostic_count, "sampling": "uniformly spaced across held-out locations over all 200 frames", "momentum_rmse": pair(np.sqrt(momentum_sse / diagnostic_count)), |
| "divergence_rmse": float(np.sqrt(div_sse / diagnostic_count)), |
| "trained_checkpoint_ad_equivalence": {"status": "passed", "max_absolute_difference": ad_max_abs}}, |
| "comparison_limits": ["Table3 exact seed/checkpoint/evaluation aggregation not specified in paper", "Paper Eq42 includes x100, but Table3 numerical entries are reported without a percent marker; both ratio and percent provided", "Native DNS covers t=0..19.9 at dt=0.1; no synthetic temporal resampling", "PyTorch port of paper TensorFlow method; coordinate AD checked against reverse-mode PyTorch"]} |
| save_json(out / "metrics.json", metrics) |
| print("FINAL METRICS " + json.dumps(metrics, ensure_ascii=False), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
