File size: 3,628 Bytes
e2703dc | 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 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | from __future__ import annotations
import sys
from pathlib import Path
from types import SimpleNamespace
import torch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
def check(name: str, tensor: torch.Tensor, expected_last_dim: int | None = None) -> None:
if not torch.isfinite(tensor).all():
raise RuntimeError(f"{name} produced non-finite values")
if expected_last_dim is not None and tensor.shape[-1] != expected_last_dim:
raise RuntimeError(
f"{name} last dim expected {expected_last_dim}, got {tensor.shape[-1]}"
)
print(f"[OK] {name}: shape={tuple(tensor.shape)}")
@torch.no_grad()
def smoke_deeponet() -> None:
from model.deeponet import DeepONetCartesianProd1D
model = DeepONetCartesianProd1D(
size=8,
in_channel_branch=1,
query_dim=1,
out_channel=1,
activation="relu",
base_model="MLP",
).eval()
x_func = torch.randn(2, 8, 1)
x_loc = torch.linspace(0.0, 1.0, 8).unsqueeze(-1)
out = model((x_func, x_loc))
check("deeponet.DeepONetCartesianProd1D", out, expected_last_dim=1)
@torch.no_grad()
def smoke_fno() -> None:
from model.fno import FNO2d
model = FNO2d(num_channels=1, modes1=4, modes2=4, width=8, initial_step=1).eval()
x = torch.randn(2, 16, 16, 1)
coords = torch.linspace(0.0, 1.0, 16)
xx, yy = torch.meshgrid(coords, coords, indexing="ij")
grid = torch.stack((xx, yy), dim=-1).unsqueeze(0).repeat(2, 1, 1, 1)
out = model(x, grid)
check("fno.FNO2d", out, expected_last_dim=1)
@torch.no_grad()
def smoke_mpnn() -> None:
from torch_geometric.data import Data
from model.mpnn import MPNN
pde = SimpleNamespace(
name="smoke",
tmin=0.0,
tmax=1.0,
resolution_t=10,
spatial_domain=[(0.0, 1.0)],
resolution=[4],
spatial_dim=1,
variables={},
)
model = MPNN(pde=pde, time_window=10, hidden_features=128, hidden_layers=1).eval()
edge_index = torch.tensor(
[[0, 1, 2, 3, 1, 2, 3, 0], [1, 2, 3, 0, 0, 1, 2, 3]],
dtype=torch.long,
)
data = Data(
x=torch.randn(4, 10, 1),
x_pos=torch.linspace(0.0, 1.0, 4).unsqueeze(-1),
t_pos=torch.zeros(4),
edge_index=edge_index,
batch=torch.zeros(4, dtype=torch.long),
variables=torch.empty(4, 0),
)
out = model(data, v=0)
check("mpnn.MPNN", out)
@torch.no_grad()
def smoke_pino_fno() -> None:
from model.pino_fno import FNO1d
model = FNO1d(
modes=[4, 4],
width=4,
layers=[4, 4, 4],
fc_dim=8,
in_dim=2,
out_dim=1,
pad_ratio=[0.0, 0.0],
).eval()
out = model(torch.randn(2, 16, 2))
check("pino_fno.FNO1d", out, expected_last_dim=1)
@torch.no_grad()
def smoke_unet() -> None:
from model.unet import UNet2d
model = UNet2d(in_channels=1, out_channels=1, init_features=2).eval()
out = model(torch.randn(1, 1, 16, 16))
check("unet.UNet2d", out)
@torch.no_grad()
def smoke_uno() -> None:
from model.uno import UNO1d
model = UNO1d(num_channels=1, width=4, initial_step=1).eval()
x = torch.randn(1, 64, 1)
grid = torch.linspace(0.0, 1.0, 64).view(1, 64, 1)
out = model(x, grid)
check("uno.UNO1d", out, expected_last_dim=1)
def main() -> int:
torch.manual_seed(0)
smoke_deeponet()
smoke_fno()
smoke_mpnn()
smoke_pino_fno()
smoke_unet()
smoke_uno()
return 0
if __name__ == "__main__":
raise SystemExit(main())
|