added the time dependent deeponet model

models/base.py

@@ -0,0 +1,74 @@
+import pytorch_lightning as pl
+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()
+
+    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 _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)
+
+        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

models/deriv_calc.py

@@ -0,0 +1,110 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import warnings
+from itertools import product
+from typing import Dict
+
+
+def gauss_pt_eval(tensor: torch.Tensor, kernels: nn.ParameterList, stride: int = 1) -> torch.Tensor:
+    if not kernels:
+        raise ValueError("No Gauss kernels provided.")
+    conv = F.conv2d
+    B, C = tensor.shape[0], tensor.shape[1]
+    device = tensor.device
+    # determine output spatial shape
+    with torch.no_grad():
+        sample_out = conv(tensor[:, :1], kernels[0].to(device), stride=stride)
+        out_spatial = sample_out.shape[2:]
+
+    results = []
+    for k in kernels:
+        k = k.to(device)
+        # apply convolution per channel
+        out_ch = [conv(tensor[:, i:i+1], k, stride=stride) for i in range(C)]
+        results.append(torch.cat(out_ch, dim=1).unsqueeze(1))
+
+    out = torch.cat(results, dim=1)
+    expected = (B, len(kernels), C) + out_spatial
+    if out.shape != expected:
+        warnings.warn(f"Shape mismatch in gauss_pt_eval: {out.shape} != {expected}")
+    return out
+
+
+class FEM2D(nn.Module):
+    """
+    Builds 2D FEM convolution kernels and evaluates derivatives.
+    """
+    def __init__(
+        self,
+        height: int,
+        width: int,
+        domain_length_x: float,
+        domain_length_y: float,
+        device: torch.device
+    ):
+        super().__init__()
+        self.height, self.width = height, width
+        self.device = device
+        # 2-point Gauss quadrature
+        self.gpx = [-0.57735, 0.57735]
+        self.kernels_dx = nn.ParameterList()
+        self.kernels_dy = nn.ParameterList()
+        self._build_kernels(domain_length_x, domain_length_y)
+
+    def _build_kernels(self, Lx: float, Ly: float):
+        hx = Lx / (self.width - 1)
+        hy = Ly / (self.height - 1)
+        # linear basis functions on [-1,1]
+        bf = lambda x: [0.5 * (1 - x), 0.5 * (1 + x)]
+        dbf = lambda x: [-0.5, 0.5]
+
+        for gx, gy in product(self.gpx, repeat=2):
+            dx = torch.zeros(2, 2, device=self.device)
+            dy = torch.zeros(2, 2, device=self.device)
+            for i, bf_x in enumerate(bf(gx)):
+                for j, bf_y in enumerate(bf(gy)):
+                    dx[j, i] = dbf(gx)[i] * (2 / hx) * bf_y
+                    dy[j, i] = bf_x * (dbf(gy)[j] * (2 / hy))
+            # store kernels with shape [1,1,2,2]
+            self.kernels_dx.append(nn.Parameter(dx.unsqueeze(0).unsqueeze(0), requires_grad=False))
+            self.kernels_dy.append(nn.Parameter(dy.unsqueeze(0).unsqueeze(0), requires_grad=False))
+
+    def eval_derivative_x(self, tensor: torch.Tensor) -> torch.Tensor:
+        return gauss_pt_eval(tensor, self.kernels_dx)
+
+    def eval_derivative_y(self, tensor: torch.Tensor) -> torch.Tensor:
+        return gauss_pt_eval(tensor, self.kernels_dy)
+
+
+class DerivativeCalculator(nn.Module):
+    """
+    Computes first spatial derivatives for 'u' and 'v' channels.
+    """
+    def __init__(
+        self,
+        height: int,
+        width: int,
+        domain_length_x: float,
+        domain_length_y: float,
+        device: torch.device,
+        channels: int = 2  # number of channels: 2 for (u,v)
+    ):
+        super().__init__()
+        self.channels = channels
+        self.fem = FEM2D(height, width, domain_length_x, domain_length_y, device)
+
+    def calculate_first_derivatives(self, y_spatial: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """
+        y_spatial: [B, C, H, W] tensor where C == channels
+        Returns a dict with keys 'u_x','u_y','v_x','v_y'.
+        """
+        deriv = {}
+        names = ['u', 'v'][:self.channels]
+        for idx, name in enumerate(names):
+            field = y_spatial[:, idx:idx+1]
+            deriv[f'{name}_x'] = self.fem.eval_derivative_x(field)
+            deriv[f'{name}_y'] = self.fem.eval_derivative_y(field)
+        return deriv
+
+    forward = calculate_first_derivatives

models/geometric_deeponet/geometric_deeponet.py

@@ -0,0 +1,36 @@
+import torch
+from models.base import BaseLightningModule
+from models.geometric_deeponet.network import GeoDeepONetTime as _GeoDeepONetTime
+from models.deriv_calc import DerivativeCalculator
+
+class GeometricDeepONetTime(BaseLightningModule):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.save_hyperparameters()
+        eff = self.hparams.output_channels - 1
+        
+        self.deriv_calc = DerivativeCalculator(
+            height=self.hparams.height,
+            width=self.hparams.width,
+            domain_length_x=self.hparams.domain_length_x,
+            domain_length_y=self.hparams.domain_length_y,
+            device=torch.device('cpu'),
+            channels=eff
+        )
+        
+        # build network args
+        net_args = {k: getattr(self.hparams, k) for k in [
+            'height', 'width', 'num_input_timesteps', 'input_channels_loc',
+            'modes', 'branch_stage1_layers', 'trunk_stage1_layers',
+            'branch_stage2_layers', 'trunk_stage2_layers',
+            'cnn_c1', 'cnn_c2', 'cnn_c3', 'cnn_fc1', 'cnn_fc2'
+        ]}
+        net_args['effective_output_channels'] = eff
+        self.model = _GeoDeepONetTime(**net_args)
+        
+    def forward(self, inputs: tuple):
+        """
+        This is the LightningModule entry point for inference.
+        It simply calls through to the underlying torch.nn.Module.
+        """
+        return self.model(inputs)        

models/geometric_deeponet/network.py

@@ -0,0 +1,152 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+import warnings
+from typing import List
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size,
+                              stride=stride, padding=padding, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+    def forward(self, x):
+        return self.relu(self.bn(self.conv(x)))
+
+class InceptionStyleCNNEncoder(nn.Module):
+    def __init__(self, input_channels: int, c1: int, c2: int, c3: int, fc1: int, fc2: int):
+        super().__init__()
+        # Branch 1
+        self.b1_conv1 = ConvBlock(input_channels, c1, 1)
+        self.b1_pool1 = nn.MaxPool2d(2, 2)
+        self.b1_conv2 = ConvBlock(c1, c2, 1)
+        self.b1_pool2 = nn.MaxPool2d(2, 2)
+        self.b1_conv3 = ConvBlock(c2, c3, 1)
+        self.b1_pool3 = nn.MaxPool2d(2, 2)
+        # Branch 2
+        self.b2_conv1a = ConvBlock(input_channels, c1, 1)
+        self.b2_conv1b = ConvBlock(c1, c1, 3, padding=1)
+        self.b2_pool1  = nn.MaxPool2d(2, 2)
+        self.b2_conv2a = ConvBlock(c1, c2, 1)
+        self.b2_conv2b = ConvBlock(c2, c2, 3, padding=1)
+        self.b2_pool2  = nn.MaxPool2d(2, 2)
+        self.b2_conv3a = ConvBlock(c2, c3, 1)
+        self.b2_conv3b = ConvBlock(c3, c3, 3, padding=1)
+        self.b2_pool3  = nn.MaxPool2d(2, 2)
+        # Branch 3
+        self.b3_conv1a = ConvBlock(input_channels, c1, 1)
+        self.b3_conv1b = ConvBlock(c1, c1, 5, padding=2)
+        self.b3_pool1  = nn.MaxPool2d(2, 2)
+        self.b3_conv2a = ConvBlock(c1, c2, 1)
+        self.b3_conv2b = ConvBlock(c2, c2, 5, padding=2)
+        self.b3_pool2  = nn.MaxPool2d(2, 2)
+        self.b3_conv3a = ConvBlock(c2, c3, 1)
+        self.b3_conv3b = ConvBlock(c3, c3, 5, padding=2)
+        self.b3_pool3  = nn.MaxPool2d(2, 2)
+        # Fusion
+        concat_channels = 3 * c3
+        self.fusion_conv1 = ConvBlock(concat_channels, fc1, 1)
+        self.fusion_pool1 = nn.MaxPool2d(2, 2)
+        self.fusion_conv2 = ConvBlock(fc1, fc2, 1)
+        self.fusion_pool2 = nn.MaxPool2d(2, 2)
+        self.flatten = nn.Flatten()
+        self.final_cnn_channels = fc2
+
+    def forward(self, x):
+        p1 = self.b1_pool3(self.b1_conv3(self.b1_pool2(self.b1_conv2(self.b1_pool1(self.b1_conv1(x))))))
+        p2 = self.b2_pool3(self.b2_conv3b(self.b2_conv3a(self.b2_pool2(self.b2_conv2b(self.b2_conv2a(self.b2_pool1(self.b2_conv1b(self.b2_conv1a(x)))))))))
+        p3 = self.b3_pool3(self.b3_conv3b(self.b3_conv3a(self.b3_pool2(self.b3_conv2b(self.b3_conv2a(self.b3_pool1(self.b3_conv1b(self.b3_conv1a(x)))))))))
+        c  = torch.cat((p1, p2, p3), dim=1)
+        f  = self.fusion_pool2(self.fusion_conv2(self.fusion_pool1(self.fusion_conv1(c))))
+        return self.flatten(f)
+
+class LinearMLP(nn.Module):
+    def __init__(self, dims: List[int], nonlin):
+        super().__init__()
+        layers = []
+        for i in range(len(dims)-1):
+            layers.append(nn.Linear(dims[i], dims[i+1]))
+            if i < len(dims)-2:
+                layers.append(nonlin())
+        self.mlp = nn.Sequential(*layers)
+    def forward(self, x):
+        return self.mlp(x)
+
+class torchSine(nn.Module):
+    def forward(self, x): return torch.sin(x)
+
+class GeoDeepONetTime(nn.Module):
+    def __init__(
+        self, height: int, width: int, num_input_timesteps: int,
+        input_channels_loc: int, effective_output_channels: int,
+        modes: int,
+        branch_stage1_layers: List[int], trunk_stage1_layers: List[int],
+        branch_stage2_layers: List[int], trunk_stage2_layers: List[int],
+        cnn_c1: int, cnn_c2: int, cnn_c3: int, cnn_fc1: int, cnn_fc2: int
+    ):
+        super().__init__()
+        if input_channels_loc != 2:
+            warnings.warn("GeoDeepONetTime expects input_channels_loc=2 (x,y). SDF will be added.")
+
+        self.input_channels_loc_base = input_channels_loc
+        self.input_channels_loc_effective = input_channels_loc + 1
+        self.effective_output_channels = effective_output_channels
+        self.modes = modes
+        self.height = height; self.width = width
+        self.num_points = height * width
+
+        # --- Branch ---
+        channels_per_step = self.effective_output_channels
+        cnn_in_ch = num_input_timesteps * channels_per_step
+        self.cnn_encoder = InceptionStyleCNNEncoder(cnn_in_ch, cnn_c1, cnn_c2, cnn_c3, cnn_fc1, cnn_fc2)
+        with torch.no_grad():
+            dummy = torch.zeros(1, cnn_in_ch, height, width)
+            flat  = self.cnn_encoder(dummy)
+            cnn_flat = flat.shape[1]
+        branch_dims1 = [cnn_flat] + branch_stage1_layers + [modes]
+        self.branch_stage_1 = LinearMLP(branch_dims1, nn.ReLU)
+        branch_dims2 = [modes] + branch_stage2_layers + [modes * effective_output_channels]
+        self.branch_stage_2 = LinearMLP(branch_dims2, nn.ReLU)
+
+        # --- Trunk ---
+        trunk_dims1 = [self.input_channels_loc_effective] + trunk_stage1_layers + [modes]
+        self.trunk_stage_1 = LinearMLP(trunk_dims1, nn.ReLU)
+        trunk_dims2 = [modes]                                 + trunk_stage2_layers + [modes * effective_output_channels]
+        self.trunk_stage_2 = LinearMLP(trunk_dims2, torchSine)
+
+        # --- Bias ---
+        self.b = nn.Parameter(torch.tensor(0.0))
+
+    def forward(self, inputs: tuple):
+        x1, _, coords, sdf = inputs[:4]
+        # --- Branch ---
+        feat = self.cnn_encoder(x1)
+        glob = self.branch_stage_1(feat)
+        # --- Trunk ---
+        # coords: [b, 2, h, w] → [b, h*w, 2]
+        c2 = rearrange(coords, 'b c h w -> b (h w) c')
+        # sdf:    [b, 1, h, w] → [b, h*w, 1]
+        sdf_flat = rearrange(sdf, 'b 1 h w -> b (h w) 1')
+        # combine into [b, h*w, 3]
+        trunk_in = torch.cat((c2, sdf_flat), dim=-1)
+        # pass each point through the trunk MLP → [b, h*w, modes]
+        local = self.trunk_stage_1(trunk_in)
+
+        # --- Merge & Stage2 --- [b, modes] → [b, 1, modes], local: [b, h*w, modes]
+        merged = glob.unsqueeze(1) * local
+        avg = merged.mean(dim=1)
+
+        out_b = self.branch_stage_2(avg)       # [b, modes*eff_out]
+        out_t = self.trunk_stage_2(merged)     # [b, h*w, modes*eff_out]
+
+        # reshape for tensor contraction
+        b_r = rearrange(out_b, 'b (m c) -> b m c', m=self.modes, c=self.effective_output_channels)
+        t_r = rearrange(out_t, 'b p (m c) -> b p m c', m=self.modes, c=self.effective_output_channels)
+
+        # compute solution and add bias
+        sol_flat = torch.einsum('bmc,bpmc->bpc', b_r, t_r) + self.b
+
+        # final shape [b, 1, p, c]
+        return rearrange(sol_flat, 'b p c -> b 1 p c')