| |
| |
| |
| |
| |
| |
| |
| |
| import os, math, pickle, argparse |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| import sys |
| BASE = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, os.path.join(BASE, 'hf_upload')) |
| from model import ISNOperator3D, get_isotropic_knn_3d_periodic |
| import physics as ph |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument('--stage', default='plot', choices=['A', 'B', 'plot']) |
| ARGS = parser.parse_args() |
|
|
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| with open(os.path.join(BASE, 'hf_upload', 'alienx_k5_stats.pkl'), 'rb') as f: |
| SD = pickle.load(f) |
| ST = SD['ST']; RE_LIST = SD['RE_LIST']; RE_MAX = SD['RE_MAX'] |
| G = SD['G']; T_FINAL = SD['T_FINAL']; N_SNAPSHOTS = SD['N_SNAPSHOTS'] |
|
|
| model = ISNOperator3D(hidden_dim=128, num_layers=4, num_heads=4, |
| re_max=RE_MAX).to(device) |
| model.G = G |
| model.load_state_dict(torch.load(os.path.join(BASE, 'hf_upload', 'alienx_k5_best.pt'), |
| map_location=device)) |
| model.eval() |
|
|
| KX, KY, KZ, K2 = ph.make_wavenumbers(G, device) |
| x = torch.arange(G, device=device, dtype=torch.float32) * (2*math.pi/G) |
| X, Y, Z = torch.meshgrid(x, x, x, indexing='ij') |
| coords = torch.stack([X.flatten(), Y.flatten(), Z.flatten()], dim=-1).unsqueeze(0) |
| knn = get_isotropic_knn_3d_periodic(G, device, dilation=max(1, G//16)) |
|
|
| def _s(v, key): |
| mu, sg = ST[key]; return (v - mu) / sg |
| def _u(v, key): |
| mu, sg = ST[key]; return v * sg + mu |
|
|
| @torch.no_grad() |
| def rollout_arrays(Re): |
| """Returns per-step arrays: mse, E_model, E_dns, Z_model, Z_dns, div.""" |
| traj = ph.generate_tgv_dataset(G=G, nu=2*math.pi/Re, T=T_FINAL, dt=0.005, |
| n_samples=N_SNAPSHOTS, device=device) |
| u = traj.to(device) |
| omega_traj = ph.velocity_to_vorticity(u, KX, KY, KZ) |
| u_curr = u[0:1].clone() |
| om_curr = omega_traj[0:1].clone() |
| U0 = ph.rms_velocity(u_curr).clamp(min=1e-6) |
| mses, Em, Ed, Zm, Zd, dv = [], [], [], [], [], [] |
| for step in range(u.shape[0] - 1): |
| om_nd = om_curr / U0 |
| u_nd = u_curr / U0 |
| k_, gm, gv, Hs = ph.scalar_derivatives(om_nd.norm(dim=-1), KX, KY, KZ) |
| re_t = torch.tensor([float(Re)], device=device) |
| pred = model(coords, _s(k_, 'k'), _s(gv, 'gv'), _s(gm, 'gm'), _s(Hs, 'H'), |
| _s(om_nd.reshape(1, -1, 3), 'om'), |
| _s(u_nd.reshape(1, -1, 3), 'u'), re_t, knn, G=G) |
| om_next = _u(pred.reshape(1, -1, 3), 'tgt').reshape(1, G, G, G, 3) * U0 |
| u_curr = ph.vorticity_to_velocity(om_next, KX, KY, KZ, K2) |
| om_curr = om_next |
| u_true = u[step+1:step+2] |
| mses.append(F.mse_loss(om_curr, omega_traj[step+1:step+2]).item()) |
| Em.append(ph.kinetic_energy(u_curr).item()) |
| Ed.append(ph.kinetic_energy(u_true).item()) |
| Zm.append(ph.enstrophy(u_curr, KX, KY, KZ).item()) |
| Zd.append(ph.enstrophy(u_true, KX, KY, KZ).item()) |
| dv.append(ph.divergence_spectral(u_curr, KX, KY, KZ).abs().mean().item()) |
| if device.type == 'cuda': |
| torch.cuda.empty_cache() |
| return {k: np.array(v) for k, v in |
| dict(mse=mses, Em=Em, Ed=Ed, Zm=Zm, Zd=Zd, div=dv).items()} |
|
|
| if ARGS.stage in ('A', 'B'): |
| todo = [300, 628] if ARGS.stage == 'A' else [900, 1200] |
| for Re in todo: |
| out = os.path.join(BASE, f'rollout_arrays_{Re}.npz') |
| if os.path.exists(out): |
| print(f'Re={Re}: cached') |
| continue |
| print(f'Re={Re}: rolling out 73 steps ...') |
| arrs = rollout_arrays(Re) |
| np.savez(out, **arrs) |
| print(f' saved {out} (final MSE {arrs["mse"][-1]:.5f}, max div {arrs["div"].max():.2e})') |
| raise SystemExit |
|
|
| |
| import matplotlib |
| matplotlib.use('Agg') |
| import matplotlib.pyplot as plt |
| from matplotlib import rcParams |
|
|
| rcParams.update({ |
| 'font.family': 'serif', |
| 'font.size': 11, |
| 'axes.grid': True, |
| 'grid.alpha': 0.3, |
| 'axes.spines.top': False, |
| 'axes.spines.right': False, |
| }) |
|
|
| COLORS = {300: '#2c7bb6', 628: '#abd9e9', 900: '#f46d43', 1200: '#d7191c'} |
| LABELS = {300: 'Re=300 (trained)', 628: 'Re=628 (trained)', |
| 900: 'Re=900 (held-out)', 1200: 'Re=1200 (trained)'} |
| DT = T_FINAL / (N_SNAPSHOTS - 1) |
| t = np.arange(1, N_SNAPSHOTS) * DT |
|
|
| data = {Re: dict(np.load(os.path.join(BASE, f'rollout_arrays_{Re}.npz'))) |
| for Re in [300, 628, 900, 1200]} |
|
|
| |
| fig, ax = plt.subplots(figsize=(6.5, 4.2)) |
| for Re in [300, 628, 900, 1200]: |
| m = data[Re]['mse'] |
| ax.plot(t, np.maximum(m, 1e-8), color=COLORS[Re], lw=1.8, label=LABELS[Re]) |
| ax.set_yscale('log') |
| ax.set_xlabel('physical time t') |
| ax.set_ylabel('pointwise MSE (model vs DNS, vorticity)') |
| ax.set_title('Autoregressive rollout error over the full trajectory (73 steps)', fontsize=11) |
| ax.legend(frameon=False, fontsize=9) |
| fig.tight_layout() |
| fig.savefig(os.path.join(BASE, 'fig1_mse_growth.png'), dpi=300) |
| plt.close(fig) |
|
|
| |
| fig, axes = plt.subplots(1, 2, figsize=(10.5, 4.2)) |
| for ax, (km, kd, ylab, title) in zip( |
| axes, |
| [('Em', 'Ed', 'kinetic energy E(t)', 'Kinetic energy'), |
| ('Zm', 'Zd', 'enstrophy Z(t)', 'Enstrophy')]): |
| for Re in [300, 628, 900, 1200]: |
| ax.plot(t, data[Re][kd], color=COLORS[Re], lw=2.4, alpha=0.9) |
| ax.plot(t, data[Re][km], color=COLORS[Re], lw=1.3, linestyle='--') |
| ax.set_xlabel('physical time t') |
| ax.set_ylabel(ylab) |
| ax.set_title(title, fontsize=11) |
| from matplotlib.lines import Line2D |
| legend = [Line2D([0], [0], color=COLORS[r], lw=2.4, label=LABELS[r]) for r in [300, 628, 900, 1200]] |
| legend += [Line2D([0], [0], color='gray', lw=2.4, label='DNS (solver)'), |
| Line2D([0], [0], color='gray', lw=1.3, linestyle='--', label='AlienX 3D (model)')] |
| fig.legend(handles=legend, loc='lower center', ncol=3, frameon=False, fontsize=9) |
| fig.tight_layout(rect=[0, 0.08, 1, 1]) |
| fig.savefig(os.path.join(BASE, 'fig2_invariants.png'), dpi=300) |
| plt.close(fig) |
|
|
| |
| fig, ax = plt.subplots(figsize=(6.5, 4.2)) |
| for Re in [300, 628, 900, 1200]: |
| ax.plot(t, np.maximum(data[Re]['div'], 1e-16), color=COLORS[Re], lw=1.8, label=LABELS[Re]) |
| ax.set_yscale('log') |
| ax.set_ylim(1e-8, 1e-4) |
| ax.set_xlabel('physical time t') |
| ax.set_ylabel(r'$|\nabla \cdot u|$ (mean abs)') |
| ax.set_title('Divergence during rollout — free by construction (Biot–Savart projection)', fontsize=11) |
| ax.legend(frameon=False, fontsize=9) |
| fig.tight_layout() |
| fig.savefig(os.path.join(BASE, 'fig3_divergence.png'), dpi=300) |
| plt.close(fig) |
|
|
| print('figures written: fig1_mse_growth.png, fig2_invariants.png, fig3_divergence.png') |
|
|
