File size: 5,993 Bytes
bdce880 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """CPU-only actual-vs-predicted field plots for OOD cases (uses best.pt).
Renders, for a few held-out (OOD) cases, a z~0 mid-plane slice colored by the
true CFD field, the model prediction, and the absolute error -- for both velocity
magnitude |U| and pressure p. Pure CPU; safe to run alongside GPU training.
Run: python -m ezflow_v3.gnn.plot_ood --n 2
"""
from __future__ import annotations
import argparse, os, sys
import numpy as np
import torch
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.tri as mtri
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("OMP_NUM_THREADS", "2")
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from ezflow_v3.gnn.etl import CaseDatasetV2
from ezflow_v3.gnn.model_v5 import MeshGraphNetV5
from ezflow_v3.gnn import features as F
from ezflow_v3.gnn.train_v5 import split
def r2(t, p):
ss = ((t - p) ** 2).sum(); tot = ((t - t.mean()) ** 2).sum() + 1e-30
return float(1 - ss / tot)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--run", default=os.path.join(os.path.dirname(__file__), "_runs", "rans_v5"))
ap.add_argument("--cache", default=r"C:\dev\ezflow_eval\cache_v3")
ap.add_argument("--out", default=r"C:\dev\ezflow_eval\ood_plots")
ap.add_argument("--ckpt", default="best.pt")
ap.add_argument("--n", type=int, default=2)
a = ap.parse_args()
os.makedirs(a.out, exist_ok=True)
ck = torch.load(os.path.join(a.run, a.ckpt), map_location="cpu", weights_only=False)
model = MeshGraphNetV5(F.NODE_FEATURE_DIM, 4, F.GLOBAL_DIM,
hidden=int(ck["args"]["hidden"]), K=int(ck["args"]["K"]),
out_dim=F.TARGET_DIM)
model.load_state_dict(ck["model"]); model.eval()
ep = ck.get("epoch", "?")
print(f"loaded {a.ckpt} (epoch {ep}, val_loss {ck.get('val_loss','?')})", flush=True)
nz = np.load(os.path.join(a.run, "norms.npz"))
ym = torch.tensor(nz["y_mean"]); ys = torch.tensor(nz["y_std"])
gm = torch.tensor(nz["g_mean"]); gs = torch.tensor(nz["g_std"])
ds = CaseDatasetV2(a.cache)
_, _, ood = split(ds)
for setname, graphs in ood.items():
sel = graphs[:a.n]
if not sel:
continue
fig, axes = plt.subplots(len(sel), 6, figsize=(23, 3.6 * len(sel) + 0.8), squeeze=False)
for r, g in enumerate(sel):
gg = g.clone(); gg.global_feat = (gg.global_feat - gm) / gs
with torch.no_grad():
pred = model(gg).numpy() * ys.numpy() + ym.numpy() # -> physical
gt = g.y.numpy()
pos = g.pos.numpy()
uT = np.linalg.norm(gt[:, :3], axis=1); uP = np.linalg.norm(pred[:, :3], axis=1)
pT, pP = gt[:, 3], pred[:, 3]
r2u, r2p = r2(uT, uP), r2(pT, pP) # R2 on the full graph
Re = float(getattr(g, "Re", 0.0))
# z~0 mid-plane slice (flow is in x-y), then crop to body + wake
z = pos[:, 2]; tol = 0.15
while (np.abs(z) < tol).sum() < 800 and tol < 1.0:
tol += 0.08
body = uT < 0.2
cx = np.median(pos[body, 0]) if body.any() else 0.0
cy = np.median(pos[body, 1]) if body.any() else 0.0
m = (np.abs(z) < tol) & (pos[:, 0] > cx - 2.5) & (pos[:, 0] < cx + 8.0) \
& (np.abs(pos[:, 1] - cy) < 3.0)
X, Y = pos[m, 0], pos[m, 1]
# triangulate the slice, mask triangles bridging the body/far gaps
tri = mtri.Triangulation(X, Y)
t = tri.triangles
e = np.concatenate([np.hypot(X[t[:, i]] - X[t[:, j]], Y[t[:, i]] - Y[t[:, j]])
for i, j in ((0, 1), (1, 2), (2, 0))])
maxe = np.maximum.reduce([np.hypot(X[t[:, 0]] - X[t[:, 1]], Y[t[:, 0]] - Y[t[:, 1]]),
np.hypot(X[t[:, 1]] - X[t[:, 2]], Y[t[:, 1]] - Y[t[:, 2]]),
np.hypot(X[t[:, 2]] - X[t[:, 0]], Y[t[:, 2]] - Y[t[:, 0]])])
tri.set_mask(maxe > 4.5 * np.median(e))
xlo, xhi = np.percentile(X, 1) - 0.3, np.percentile(X, 99) + 0.3
ylo, yhi = np.percentile(Y, 1) - 0.3, np.percentile(Y, 99) + 0.3
def panel(c, val, title, vmin, vmax, cmap):
lv = np.linspace(vmin, vmax, 25)
cs = axes[r, c].tricontourf(tri, val[m], levels=lv, cmap=cmap, extend="both")
axes[r, c].set_title(title, fontsize=9)
axes[r, c].set_aspect("equal"); axes[r, c].set_xticks([]); axes[r, c].set_yticks([])
axes[r, c].set_xlim(xlo, xhi); axes[r, c].set_ylim(ylo, yhi)
plt.colorbar(cs, ax=axes[r, c], fraction=0.046, pad=0.02)
vU = (float(min(uT.min(), uP.min())), float(max(uT.max(), uP.max())))
vP = (float(min(pT.min(), pP.min())), float(max(pT.max(), pP.max())))
eU = np.abs(uP - uT); eP = np.abs(pP - pT)
panel(0, uT, f"{g.cid}\nRe={Re:.0f} |U| TRUE", *vU, "viridis")
panel(1, uP, f"|U| PRED (R2={r2u:.3f})", *vU, "viridis")
panel(2, eU, "|U| |error|", 0.0, float(eU[m].max() + 1e-9), "magma")
panel(3, pT, "p TRUE", *vP, "coolwarm")
panel(4, pP, f"p PRED (R2={r2p:.3f})", *vP, "coolwarm")
panel(5, eP, "p |error|", 0.0, float(eP[m].max() + 1e-9), "magma")
print(f" {setname:12s} {g.cid:24s} Re={Re:7.0f} |U|R2={r2u:.3f} pR2={r2p:.3f} (slice n={int(m.sum())}, tol={tol:.2f})", flush=True)
fig.suptitle(f"OOD: {setname} - actual vs predicted ({a.ckpt} ep{ep}, z~0 slice)", fontsize=13)
fig.tight_layout(rect=[0, 0, 1, 0.96])
outp = os.path.join(a.out, f"ood_{setname}.png")
fig.savefig(outp, dpi=110); plt.close(fig)
print(f"saved {outp}", flush=True)
print("PLOT_OOD_DONE", flush=True)
if __name__ == "__main__":
main()
|