| import torch |
| import numpy as np |
| import matplotlib.pyplot as plt |
| import matplotlib.animation as animation |
| from mpl_toolkits.axes_grid1 import make_axes_locatable |
| from scipy.spatial import cKDTree |
| import os |
|
|
| from dataset import CFDReconstructionDataset |
| from architectures import PIGU_Hybrid |
| from model import NavierStokesURANS |
|
|
| |
| |
| |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
| SENSOR_COUNT = 65 |
| DT = 0.05 |
| FRAMES_TO_RENDER = 60 |
|
|
| |
| MODEL_WEIGHTS_PATH = "results_ablation_v5/model_ep400.pth" |
| OUTPUT_FILENAME = "eval_results/var5_animation_with_metrics.gif" |
|
|
| PATH_UNSTEADY = "../../Ablation Experiment/dataset_sin/flow_one_period.npy" |
| PATH_MEAN = "../../Ablation Experiment/dataset_sin/mean_flow_steady.npy" |
|
|
| VAR_IDX = 0 |
| VAR_NAME = ['u-velocity', 'v-velocity', 'Pressure'][VAR_IDX] |
|
|
| def main(): |
| os.makedirs(os.path.dirname(OUTPUT_FILENAME), exist_ok=True) |
| |
| print(f"[*] Initializing CFD Animation & Evaluation for Variant 5...") |
| dataset = CFDReconstructionDataset(PATH_UNSTEADY, PATH_MEAN, SENSOR_COUNT, dt=DT) |
| |
| stats_max = dataset.stats['max'].to(DEVICE) |
| stats_min = dataset.stats['min'].to(DEVICE) |
| |
| dx_val = dataset.box_len[0] / (dataset.W - 1) |
| dy_val = dataset.box_len[1] / (dataset.H - 1) |
| dx_tensor = torch.tensor(dx_val, device=DEVICE) |
| dy_tensor = torch.tensor(dy_val, device=DEVICE) |
| |
| def denormalize_batch(norm_tensor): |
| return (norm_tensor + 1) / 2 * (stats_max - stats_min) + stats_min |
|
|
| print(f"[*] Loading Model Weights from: {MODEL_WEIGHTS_PATH}") |
| model = PIGU_Hybrid(sensor_in_dim=3, sensor_count=SENSOR_COUNT).to(DEVICE) |
| |
| if not os.path.exists(MODEL_WEIGHTS_PATH): |
| print(f"[!] Warning: Weights {MODEL_WEIGHTS_PATH} not found. Please train V5 first.") |
| return |
| |
| model.load_state_dict(torch.load(MODEL_WEIGHTS_PATH, map_location=DEVICE)) |
| model.eval() |
| |
| pde_engine = NavierStokesURANS(dataset.stats).to(DEVICE) |
|
|
| |
| sensor_phys_x = [dataset.grid_X[p[0], p[1]] for p in dataset.sensor_indices] |
| sensor_phys_y = [dataset.grid_Y[p[0], p[1]] for p in dataset.sensor_indices] |
|
|
| print("[*] Building Geometry Mask...") |
| raw_mean = np.load(PATH_MEAN).astype(np.float32) |
| coords_mean = raw_mean[:, :2] |
| tree = cKDTree(coords_mean) |
| |
| grid_pts = np.stack([dataset.grid_X.ravel(), dataset.grid_Y.ravel()], axis=-1) |
| dist, _ = tree.query(grid_pts) |
| threshold = max(dx_val, dy_val) * 2.5 |
| is_solid_wall = (dist > threshold).reshape(dataset.H, dataset.W) |
| is_solid_wall_float = is_solid_wall.astype(float) |
| |
| valid_mask_torch = ~torch.from_numpy(is_solid_wall).to(DEVICE) |
|
|
| print(f"[*] Running Inference and Calculating Metrics for {FRAMES_TO_RENDER} frames...") |
| true_frames, pred_frames, err_frames = [], [], [] |
| |
| total_l2 = 0.0 |
| total_rmse = 0.0 |
| total_resc = 0.0 |
| num_frames = min(FRAMES_TO_RENDER, len(dataset)) |
|
|
| with torch.no_grad(): |
| for idx in range(num_frames): |
| s_val_t, s_pos, grid_pos_norm, s_val_next, mean_flow = dataset[idx] |
| true_norm_t = dataset.data[idx] |
| |
| s_val_t = s_val_t.unsqueeze(0).to(DEVICE) |
| s_pos = s_pos.unsqueeze(0).to(DEVICE) |
| grid_pos_norm = grid_pos_norm.unsqueeze(0).to(DEVICE) |
| s_val_next = s_val_next.unsqueeze(0).to(DEVICE) |
| mean_flow = mean_flow.unsqueeze(0).to(DEVICE) |
| true_norm_t = true_norm_t.unsqueeze(0).to(DEVICE) |
|
|
| pred_norm_t = model(s_val_t, s_pos, grid_pos_norm, base_flow=mean_flow) |
| pred_norm_next = model(s_val_next, s_pos, grid_pos_norm, base_flow=mean_flow) |
| |
| pred_phys_t = denormalize_batch(pred_norm_t) |
| true_phys_t = denormalize_batch(true_norm_t) |
| pred_phys_next = denormalize_batch(pred_norm_next) |
| |
| p_var_torch = pred_phys_t[0, VAR_IDX] |
| t_var_torch = true_phys_t[0, VAR_IDX] |
| |
| diff_valid = (p_var_torch - t_var_torch)[valid_mask_torch] |
| t_valid = t_var_torch[valid_mask_torch] |
| |
| frame_l2 = torch.norm(diff_valid) / (torch.norm(t_valid) + 1e-8) |
| frame_rmse = torch.sqrt(torch.mean(diff_valid**2)) |
| |
| _, _, res_c = pde_engine(pred_phys_t, pred_phys_next, None, dx=dx_tensor, dy=dy_tensor) |
| res_c_valid = res_c[0, 0][valid_mask_torch] |
| frame_resc = torch.sqrt(torch.mean(res_c_valid**2)) |
| |
| total_l2 += frame_l2.item() |
| total_rmse += frame_rmse.item() |
| total_resc += frame_resc.item() |
| |
| p_var = p_var_torch.cpu().numpy() |
| t_var = t_var_torch.cpu().numpy() |
| |
| p_var[is_solid_wall] = np.nan |
| t_var[is_solid_wall] = np.nan |
| err_var = np.abs(p_var - t_var) |
| |
| true_frames.append(t_var) |
| pred_frames.append(p_var) |
| err_frames.append(err_var) |
|
|
| avg_l2 = total_l2 / num_frames |
| avg_rmse = total_rmse / num_frames |
| avg_resc = total_resc / num_frames |
| |
| print(f"\n[Metrics] L2 Rel Error: {avg_l2:.4f} | RMSE: {avg_rmse:.4f} | Cont. Res: {avg_resc:.2e}") |
|
|
| |
| |
| |
| print(f"[*] Generating Professional CFD-Style Animation...") |
| |
| vmin = np.nanmin(true_frames) |
| vmax = np.nanmax(true_frames) |
| err_vmax = np.nanmax(err_frames) * 0.8 |
|
|
| fig, axes = plt.subplots(1, 3, figsize=(18, 5.5), facecolor='white') |
| fig.subplots_adjust(bottom=0.2) |
| fig.suptitle(f'Variant 5 (w/o Base Flow) Reconstruction - {VAR_NAME}', fontsize=18, fontweight='bold', y=0.98) |
| |
| extent = [dataset.grid_X.min(), dataset.grid_X.max(), dataset.grid_Y.min(), dataset.grid_Y.max()] |
|
|
| for ax in axes: |
| ax.set_facecolor('#c0c0c0') |
| ax.set_aspect('equal') |
| ax.set_xlabel("x") |
| ax.set_ylabel("y") |
| ax.spines['top'].set_visible(False) |
| ax.spines['right'].set_visible(False) |
| |
| im_true = axes[0].imshow(true_frames[0], cmap='turbo', origin='lower', extent=extent, vmin=vmin, vmax=vmax, interpolation='bicubic') |
| axes[0].contour(is_solid_wall_float, levels=[0.5], extent=extent, colors='black', linewidths=1.5, origin='lower') |
| axes[0].set_title("CFD Ground Truth", fontsize=14, pad=10) |
| |
| im_pred = axes[1].imshow(pred_frames[0], cmap='turbo', origin='lower', extent=extent, vmin=vmin, vmax=vmax, interpolation='bicubic') |
| axes[1].contour(is_solid_wall_float, levels=[0.5], extent=extent, colors='black', linewidths=1.5, origin='lower') |
| axes[1].set_title("Direct Full-Flow Prediction", fontsize=14, pad=10) |
| axes[1].scatter(sensor_phys_x, sensor_phys_y, c='white', edgecolors='black', s=18, alpha=0.9, zorder=5) |
| |
| divider = make_axes_locatable(axes[1]) |
| cax = divider.append_axes("right", size="5%", pad=0.1) |
| fig.colorbar(im_pred, cax=cax, label=f'{VAR_NAME} Magnitude') |
| |
| im_err = axes[2].imshow(err_frames[0], cmap='inferno', origin='lower', extent=extent, vmin=0, vmax=err_vmax, interpolation='bicubic') |
| axes[2].contour(is_solid_wall_float, levels=[0.5], extent=extent, colors='white', linewidths=1.5, origin='lower') |
| axes[2].set_title("Absolute Error", fontsize=14, pad=10) |
| |
| divider_err = make_axes_locatable(axes[2]) |
| cax_err = divider_err.append_axes("right", size="5%", pad=0.1) |
| fig.colorbar(im_err, cax=cax_err, label='Error Magnitude') |
|
|
| metrics_text = ( |
| f"Quantitative Evaluation ({FRAMES_TO_RENDER} frames average)\n" |
| f"$L_2$ Relative Error: {avg_l2:.4f} | RMSE: {avg_rmse:.4f} | Continuity Residual: {avg_resc:.2e}" |
| ) |
| fig.text(0.5, 0.05, metrics_text, ha='center', va='center', fontsize=14, fontweight='bold', |
| bbox=dict(facecolor='#f8f9fa', alpha=0.9, edgecolor='#adb5bd', boxstyle='round,pad=0.5')) |
|
|
| def update(frame_idx): |
| im_true.set_array(true_frames[frame_idx]) |
| im_pred.set_array(pred_frames[frame_idx]) |
| im_err.set_array(err_frames[frame_idx]) |
| return [im_true, im_pred, im_err] |
|
|
| ani = animation.FuncAnimation(fig, update, frames=len(true_frames), interval=100, blit=True) |
|
|
| if OUTPUT_FILENAME.endswith('.mp4'): |
| ani.save(OUTPUT_FILENAME, writer='ffmpeg', fps=10, dpi=200) |
| else: |
| ani.save(OUTPUT_FILENAME, writer='pillow', fps=10, dpi=150) |
| |
| print(f"[+] Animation successfully saved to: {OUTPUT_FILENAME}") |
|
|
| if __name__ == "__main__": |
| main() |
