from pathlib import Path import io import zipfile import tempfile from functools import lru_cache import numpy as np import torch import gradio as gr from huggingface_hub import hf_hub_download import matplotlib.pyplot as plt import imageio.v2 as imageio from mpl_toolkits.axes_grid1 import make_axes_locatable from einops import rearrange from models.geometric_deeponet.geometric_deeponet import GeometricDeepONetTime # ---------------- Config ---------------- REPO_ID = "arabeh/DeepONet-FlowBench-FPO" CKPTS = { "1": "checkpoints/time-dependent-deeponet_1in.ckpt", "4": "checkpoints/time-dependent-deeponet_4in.ckpt", "8": "checkpoints/time-dependent-deeponet_8in.ckpt", "16": "checkpoints/time-dependent-deeponet_16in.ckpt", } SAMPLES_DIR = Path("sample_cases") DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") TMP = Path(tempfile.gettempdir()) RANGES = { "u": (-2.0, 2.0), "v": (-1.0, 1.0), } def _tag() -> str: # unique per request (avoids filename collisions across sessions) return next(tempfile._get_candidate_names()) def _tmp(tag: str, name: str) -> str: out_dir = TMP / f"deeponet_fpo_{tag}" out_dir.mkdir(parents=True, exist_ok=True) return str(out_dir / name) # ---------------- Samples ---------------- def list_samples(): if not SAMPLES_DIR.is_dir(): return [] ids = [] for p in SAMPLES_DIR.glob("sample_*_input.npy"): # sample_{id}_input.npy sid = p.stem.split("_")[1] if sid.isdigit(): ids.append(sid) return sorted(set(ids), key=int) def load_sample(sample_id: str): sdf = np.load(SAMPLES_DIR / f"sample_{sample_id}_input.npy").astype(np.float32) # [1,H,W] y = np.load(SAMPLES_DIR / f"sample_{sample_id}_output.npy").astype(np.float32) # [T,2,H,W] return sdf, y # ---------------- Model ---------------- @lru_cache(maxsize=4) def load_model(history_s: int) -> GeometricDeepONetTime: ckpt_path = hf_hub_download(REPO_ID, CKPTS[str(history_s)]) model = GeometricDeepONetTime.load_from_checkpoint(ckpt_path, map_location=DEVICE) return model.eval().to(DEVICE) def static_tensors(hparams, sdf_np: np.ndarray): _, H, W = sdf_np.shape x = np.linspace(0.0, float(hparams.domain_length_x), W, dtype=np.float32) y = np.linspace(0.0, float(hparams.domain_length_y), H, dtype=np.float32) yv, xv = np.meshgrid(y, x, indexing="ij") coords = np.stack([xv, yv], axis=0)[None] # [1,2,H,W] sdf_t = torch.from_numpy(sdf_np)[None].to(DEVICE) # [1,1,H,W] coords_t = torch.from_numpy(coords).to(DEVICE) # [1,2,H,W] re_t = torch.zeros_like(sdf_t) # [1,1,H,W] return sdf_t, coords_t, re_t, H, W # ---------------- Rollout + metrics ---------------- def rollout(sample_id: str, history_s: str): s = int(history_s) model = load_model(s) sdf, y_true = load_sample(sample_id) T, C, H, W = y_true.shape if C != 2: raise ValueError(f"Expected 2 channels (u,v), got {C}") s = min(s, T - 1) # ensure s < T sdf_t, coords_t, re_t, _, _ = static_tensors(model.hparams, sdf) y_pred = np.zeros_like(y_true) y_pred[:s] = y_true[:s] history = y_true[:s].copy() # [s,2,H,W] for t in range(s, T): branch = rearrange(history, "nb c h w -> (nb c) h w")[None] # [1,s*2,H,W] branch_t = torch.from_numpy(branch).to(DEVICE) with torch.no_grad(): y_hat = model((branch_t, re_t, coords_t, sdf_t)) # [1,1,p,2] frame = y_hat[0, 0].view(H, W, 2).permute(2, 0, 1).cpu().numpy() # [2,H,W] y_pred[t] = frame history = frame[None] if s == 1 else np.concatenate([history[1:], frame[None]], axis=0) return y_true, y_pred, s def rollout_errors(y_true: np.ndarray, y_pred: np.ndarray, s: int): yt = y_true[s:] yp = y_pred[s:] diff = yp - yt ts = np.arange(s, y_true.shape[0]) def rel(comp: int): d = diff[:, comp].reshape(len(ts), -1) t = yt[:, comp].reshape(len(ts), -1) return np.linalg.norm(d, axis=1) / np.linalg.norm(t, axis=1) err_u = rel(0) err_v = rel(1) return ts, err_u, err_v, float(err_u.mean()), float(err_v.mean()) def pair_png(gt2d: np.ndarray, pred2d: np.ndarray, label: str, t: int) -> bytes: vmin, vmax = RANGES.get(label, (-1.0, 1.0)) # fallback if label changes fig, ax = plt.subplots(1, 2, figsize=(6.5, 2.6)) ax[0].imshow(gt2d, origin="lower", vmin=vmin, vmax=vmax) ax[0].set_title(f"{label} GT – t={t}") ax[0].axis("off") im2 = ax[1].imshow(pred2d, origin="lower", vmin=vmin, vmax=vmax) ax[1].set_title(f"{label} Pred – t={t}") ax[1].axis("off") # Colorbar height == ax[1] image height divider = make_axes_locatable(ax[1]) cax = divider.append_axes("right", size="5%", pad=0.05) fig.colorbar(im2, cax=cax) buf = io.BytesIO() fig.savefig(buf, format="png", bbox_inches="tight", dpi=110) plt.close(fig) return buf.getvalue() def write_gif(tag: str, y_true: np.ndarray, y_pred: np.ndarray, comp: int, label: str) -> str: path = _tmp(tag, f"{label}_rollout.gif") with imageio.get_writer(path, mode="I", duration=0.1, loop=0) as w: for t in range(y_true.shape[0]): png = pair_png(y_true[t, comp], y_pred[t, comp], label, t) w.append_data(imageio.imread(io.BytesIO(png))) return path def write_zip(tag: str, y_true: np.ndarray, y_pred: np.ndarray, comp: int, label: str) -> str: path = _tmp(tag, f"{label}_frames.zip") with zipfile.ZipFile(path, "w", compression=zipfile.ZIP_DEFLATED) as zf: for t in range(y_true.shape[0]): zf.writestr(f"{label}_frame_{t:03d}.png", pair_png(y_true[t, comp], y_pred[t, comp], label, t)) return path def write_error_assets(tag: str, ts: np.ndarray, err_u: np.ndarray, err_v: np.ndarray): png = _tmp(tag, "relL2_vs_time.png") csv = _tmp(tag, "relL2_vs_time.csv") np.savetxt( csv, np.c_[ts, err_u, err_v], delimiter=",", header="timestep,rel_L2_u,rel_L2_v", comments="", ) fig, ax = plt.subplots(figsize=(5, 3)) ax.plot(ts, err_u, label="u") ax.plot(ts, err_v, label="v") ax.set_xlabel("Timestep") ax.set_ylabel("Relative L2") ax.set_title("Rollout rel. L2 vs time") ax.legend() ax.grid(True, alpha=0.3) fig.savefig(png, dpi=120, bbox_inches="tight") plt.close(fig) return png, csv # ---------------- Gradio callback ---------------- def predict_rollout(sample_id: str, history_s: str): tag = _tag() y_true, y_pred, s = rollout(sample_id, history_s) ts, err_u, err_v, avg_u, avg_v = rollout_errors(y_true, y_pred, s) u_gif = write_gif(tag, y_true, y_pred, 0, "u") v_gif = write_gif(tag, y_true, y_pred, 1, "v") u_zip = write_zip(tag, y_true, y_pred, 0, "u") v_zip = write_zip(tag, y_true, y_pred, 1, "v") err_png, csv = write_error_assets(tag, ts, err_u, err_v) metrics = ( f"Rollout relative L2 error (averaged over t ≥ {s}):\n" f" u: {avg_u:.3e}\n" f" v: {avg_v:.3e}" ) return (u_gif, u_gif, u_zip, v_gif, v_gif, v_zip, err_png, csv, metrics) # ---------------- UI builder ---------------- def build_demo(): sample_choices = list_samples() or ["0"] return gr.Interface( fn=predict_rollout, inputs=[ gr.Radio(sample_choices, value=sample_choices[0], label="Sample ID"), gr.Radio(["1", "4", "8", "16"], value="16", label="History length s"), ], outputs=[ gr.Image(type="filepath", label="u rollout (GIF)"), gr.File(label="Download u rollout (GIF)"), gr.File(label="Download all u frames (ZIP)"), gr.Image(type="filepath", label="v rollout (GIF)"), gr.File(label="Download v rollout (GIF)"), gr.File(label="Download all v frames (ZIP)"), gr.Image(type="filepath", label="Relative L2 vs time"), gr.File(label="Download L2 vs time (CSV)"), gr.Textbox(label="Summary metrics"), ], title="Time-Dependent DeepONet – FPO Rollout Demo", description=( "Auto-regressive 60-step rollout of u and v fields for a selected sample. " "Choose history length s (1, 4, 8, 16). Download videos/frames and relative error vs time (CSV)." ), ) if __name__ == "__main__": build_demo().launch()