Add v1/v2 apps, update README and deps

.gitignore

@@ -0,0 +1,3 @@
+.venv/
+**/__pycache__/
+*.pyc

README.md

@@ -1,6 +1,5 @@
 ---
 title: DeepONet FPO Demo
-emoji: 🐨
 colorFrom: green
 colorTo: green
 sdk: gradio
@@ -11,4 +10,48 @@ license: mit
 short_description: 'Demo of unsteady flow around varied geometries '
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# DeepONet FPO Demo (FlowBench)
+
+This Space runs **time-dependent DeepONet** checkpoints (s ∈ {1,4,8,16}) to generate **auto-regressive rollouts** of 2D velocity fields **(u, v)** around complex geometries (FPO / FlowBench).
+
+You have two runnable apps:
+
+- **`app_v1.py` (GT + metrics)**
+  - Uses `sample_cases/` containing the **full target sequence**.
+  - Produces:
+    - GT vs Pred **GIFs** for u and v
+    - ZIPs with all frames
+    - Relative L2 vs time **plot** + **CSV**
+    - Summary metrics (avg rel-L2 over t ≥ s)
+
+- **`app_v2.py` (prediction-only, arbitrary rollout length)**
+  - Uses `sample_cases/few_timesteps/` containing **only the first 16 GT frames** (enough to seed any checkpoint).
+  - Produces:
+    - Prediction-only **GIFs** for u and v
+    - ZIPs with all frames
+    - Short run summary
+  - User chooses rollout length **N** (can be larger than 16; seeding uses only the first `s` frames).
+
+## Sample format
+
+Each sample uses two files:
+
+- `sample_{id}_input.npy` → SDF geometry: **[1, H, W]**
+- `sample_{id}_output.npy` → velocity sequence: **[T, 2, H, W]** (channels are u, v)
+
+For **v2**, `T = 16` and files must be located in:
+- `sample_cases/few_timesteps/`
+
+## Checkpoints
+
+Checkpoints are downloaded from the Hub at runtime:
+
+- `checkpoints/time-dependent-deeponet_1in.ckpt`
+- `checkpoints/time-dependent-deeponet_4in.ckpt`
+- `checkpoints/time-dependent-deeponet_8in.ckpt`
+- `checkpoints/time-dependent-deeponet_16in.ckpt`
+
+Repo ID used by both apps:
+```text
+arabeh/DeepONet-FlowBench-FPO
+

app.py

@@ -1,22 +1,12 @@
-import torch
-from huggingface_hub import hf_hub_download
-from models.geometric_deeponet.geometric_deeponet import GeometricDeepONetTime
+import os
 
-REPO_ID = "arabeh/DeepONet-FlowBench-FPO"
-CKPTS = {
-    "s=1":  "checkpoints/time-dependent-deeponet_1in.ckpt",
-    "s=4":  "checkpoints/time-dependent-deeponet_4in.ckpt",
-    "s=8":  "checkpoints/time-dependent-deeponet_8in.ckpt",
-    "s=16": "checkpoints/time-dependent-deeponet_16in.ckpt",
-}
+# default to v2; set DEMO_VERSION=v1 to run v1
+ver = os.getenv("DEMO_VERSION", "v2").lower()
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+if ver == "v1":
+    from app_v1 import demo
+else:
+    from app_v2 import demo
 
-def load_model(s_key="s=4"):
-    ckpt_path = hf_hub_download(REPO_ID, CKPTS[s_key])
-    model = GeometricDeepONetTime.load_from_checkpoint(
-        ckpt_path,
-        map_location=device,
-    )
-    model.eval()
-    return model.to(device)
+if __name__ == "__main__":
+    demo.launch()

app_v1.py

@@ -0,0 +1,251 @@
+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:
+    return str(TMP / f"{tag}_{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 ----------------
+sample_choices = list_samples() or ["0"]
+
+demo = 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__":
+    demo.launch()
+

app_v2.py

@@ -0,0 +1,213 @@
+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",
+}
+
+# v2 samples live here (only 16 GT timesteps per sample)
+SAMPLES_DIR = Path("sample_cases") / "few_timesteps"
+
+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:
+    return next(tempfile._get_candidate_names())
+
+
+def _tmp(tag: str, name: str) -> str:
+    return str(TMP / f"{tag}_{name}")
+
+
+# ---------------- Samples ----------------
+def list_samples():
+    if not SAMPLES_DIR.is_dir():
+        return []
+    ids = []
+    for p in SAMPLES_DIR.glob("sample_*_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]
+    y16 = np.load(SAMPLES_DIR / f"sample_{sample_id}_output.npy").astype(np.float32)  # [16,2,H,W]
+    return sdf, y16
+
+
+# ---------------- 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 ----------------
+def rollout_pred(sample_id: str, history_s: str, n_steps: int):
+    s = int(history_s)
+    n_steps = int(n_steps)
+
+    if n_steps <= 0:
+        raise ValueError("Number of rollout steps must be a positive integer.")
+    if n_steps < s:
+        n_steps = s  # must have at least s frames to seed
+
+    model = load_model(s)
+    sdf, y16 = load_sample(sample_id)
+
+    # Expect [16,2,H,W] (or more), but we ONLY use first s to seed the model.
+    if y16.ndim != 4 or y16.shape[1] != 2:
+        raise ValueError(f"Expected y shape [T,2,H,W], got {y16.shape}")
+    if y16.shape[0] < s:
+        raise ValueError(f"Sample only has {y16.shape[0]} timesteps, but checkpoint needs s={s}.")
+
+    _, _, H, W = y16.shape
+    sdf_t, coords_t, re_t, _, _ = static_tensors(model.hparams, sdf)
+
+    seed = y16[:s].copy()  # [s,2,H,W] (GT seed only)
+    y_out = np.zeros((n_steps, 2, H, W), dtype=np.float32)
+    y_out[:s] = seed
+
+    history = seed.copy()
+    for t in range(s, n_steps):
+        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().astype(np.float32)  # [2,H,W]
+        y_out[t] = frame
+        history = frame[None] if s == 1 else np.concatenate([history[1:], frame[None]], axis=0)
+
+    return y_out, s
+
+
+# ---------------- Rendering (prediction-only) ----------------
+def single_png(field2d: np.ndarray, label: str, t: int) -> bytes:
+    vmin, vmax = RANGES.get(label, (-1.0, 1.0))
+
+    fig, ax = plt.subplots(1, 1, figsize=(3.4, 2.8))
+    im = ax.imshow(field2d, origin="lower", vmin=vmin, vmax=vmax)
+    ax.set_title(f"{label} – t={t}")
+    ax.axis("off")
+
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="6%", pad=0.05)
+    fig.colorbar(im, cax=cax)
+
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png", bbox_inches="tight", dpi=120)
+    plt.close(fig)
+    return buf.getvalue()
+
+
+def write_gif(tag: str, y: 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.shape[0]):
+            png = single_png(y[t, comp], label, t)
+            w.append_data(imageio.imread(io.BytesIO(png)))
+    return path
+
+
+def write_zip(tag: str, y: 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.shape[0]):
+            zf.writestr(f"{label}_frame_{t:03d}.png", single_png(y[t, comp], label, t))
+    return path
+
+# ---------------- Gradio callback ----------------
+def run_v2(sample_id: str, history_s: str, n_steps: int):
+    tag = _tag()
+    y, s = rollout_pred(sample_id, history_s, n_steps)
+
+    u_gif = write_gif(tag, y, comp=0, label="u")
+    v_gif = write_gif(tag, y, comp=1, label="v")
+    u_zip = write_zip(tag, y, comp=0, label="u")
+    v_zip = write_zip(tag, y, comp=1, label="v")
+
+    summary = (
+        f"Seeded with s={s} timesteps from {SAMPLES_DIR}.\n"
+        f"Generated rollout length N={y.shape[0]} (frames labeled seed for t<s, pred for t≥s)."
+    )
+
+    return (
+        u_gif, u_gif, u_zip,
+        v_gif, v_gif, v_zip,
+        summary,
+    )
+
+
+# ---------------- UI ----------------
+sample_choices = list_samples() or ["0"]
+history_choices = ["1", "4", "8", "16"]
+
+demo = gr.Interface(
+    fn=run_v2,
+    inputs=[
+        gr.Radio(sample_choices, value=sample_choices[0], label="Sample ID"),
+        gr.Radio(history_choices, value="16", label="History length s (checkpoint)"),
+        gr.Number(value=60, precision=0, label="Rollout steps N (total frames)"),
+    ],
+    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.Textbox(label="Run summary"),
+    ],
+    title="Time-Dependent DeepONet –  FPO Rollout Demo",
+    description=(
+        "Auto-regressive rollout of u and v fields for a selected sample. "
+        "Choose history length s (1, 4, 8, 16). Download videos/frames."
+    ),
+)
+
+if __name__ == "__main__":
+    demo.launch()

models/base.py

@@ -3,72 +3,12 @@ import torch
 
 
 class BaseLightningModule(pl.LightningModule):
-    def configure_optimizers(self):
-        return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)
-
-    def _masked_mse(self, y_hat, y_true, sdf):
-        mask = (sdf > 0).flatten(1).unsqueeze(-1)
-        se = ((y_hat - y_true) ** 2) * mask
-        return se.sum() / mask.sum()
+    """
+    Minimal LightningModule base for the demo.
+    """
 
-    def training_step(self, batch, batch_idx):
-        (branch, re, coords, sdf), tgt = batch
-        y_hat = self.model((branch, re, coords, sdf))
-        if self.hparams.use_derivative_loss:
-            loss = self._derivative_loss(y_hat, tgt, sdf)
-        else:
-            loss = self._masked_mse(y_hat, tgt, sdf)        
-        
-        self.log('train_loss', loss)
-        return loss
-
-    def validation_step(self, batch, batch_idx):
-        (branch, re, coords, sdf), tgt = batch
-        y_hat = self.model((branch, re, coords, sdf))
-        if self.hparams.use_derivative_loss:
-            loss = self._derivative_loss(y_hat, tgt, sdf)
-        else:
-            loss = self._masked_mse(y_hat, tgt, sdf)                
-        
-        self.log('val_loss', loss)
-        return loss
+    def configure_optimizers(self):
 
-    def _derivative_loss(self, y_hat, y_true, sdf):
-        # --- reshape [B,1,p,C] → [B,C,H,W] ---
-        B, _, p, C = y_hat.shape
-        H, W = self.hparams.height, self.hparams.width
-        yh = y_hat.squeeze(1).permute(0,2,1).reshape(B, C, H, W)
-        yt = y_true.squeeze(1).permute(0,2,1).reshape(B, C, H, W)
+        lr = getattr(self.hparams, "lr", 1e-3)
+        return torch.optim.Adam(self.parameters(), lr=lr)
 
-        deriv_hat = self.deriv_calc(yh)
-        deriv_true = self.deriv_calc(yt)
-        fluid_mask = (sdf > 0)  # [B,1,H,W]
-        delta = self.hparams.domain_length_y / H
-        loss = 0.0  
-        # Derivative tensors come out at resolution (H-1)x(W-1) so crop the fluid_mask to match:
-        dm = fluid_mask[:, :, :-1, :-1].unsqueeze(1)  # → [B,1,1,H-1,W-1]
-        for key in ('u_x','u_y','v_x','v_y'):
-            diff = deriv_hat[key] - deriv_true[key]     # [B,ngp,1,H-1,W-1]
-            # apply mask before averaging
-            deriv_loss  = delta * (diff.pow(2) * dm).sum() / dm.sum()
-            self.log(f"deriv_loss/{key}", deriv_loss, on_step=False, on_epoch=True)
-            loss = loss + deriv_loss
-            
-        inner = (sdf > 0) & (sdf <= delta)  # [B,1,H,W]
-        if inner.any().item():
-            u_hat = yh[:, 0:1]  # [B,1,H,W]
-            v_hat = yh[:, 1:2]
-            if self.hparams.use_zero_bc:
-                bc_loss = 1000 * (u_hat[inner].pow(2) + v_hat[inner].pow(2)).mean()
-                             
-            else:
-                u_true = yt[:, 0:1]
-                v_true = yt[:, 1:2]
-                u_target = u_true[inner]
-                v_target = v_true[inner]
-                bc_loss = ((u_hat[inner] - u_target).pow(2) + (v_hat[inner] - v_target).pow(2)).mean()
-            
-            self.log("boundary_bc_loss", bc_loss, on_step=False, on_epoch=True)
-            loss = loss + bc_loss
-                      
-        return loss

requirements.txt

@@ -2,6 +2,7 @@ gradio
 torch
 pytorch-lightning
 huggingface_hub
-einops
 numpy
-matplotlib
+imageio
+einops
+matplotlib

sample_cases/few_timesteps/sample_0_input.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:698f0a4c56d4d2beb1c4e9b3ccaba1f9c104a69ec60ca9501470589ff68f69f1
+size 1048704

sample_cases/few_timesteps/sample_0_output.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:12f2b8972a72ecd5a74294160f6d2fe38c2dc341364f85f25e235b4808168297
+size 67108992

sample_cases/few_timesteps/sample_1_input.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7af7b5ee1e43eb97b1d99fb80a012b9edc4607d5b3139afc61c30d9aca4c60ed
+size 1048704

sample_cases/few_timesteps/sample_1_output.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:56650a750acd232d9fb173b6860c45d080b66ff3a239c13b6957e91674e12dee
+size 67108992