add module and first model

ThermalSolver/model/local_gno.py

@@ -0,0 +1,108 @@
+import torch
+from torch import nn
+from torch_geometric.nn import MessagePassing
+
+
+# ---- FiLM that starts as identity and normalizes the target ----
+class FiLM(nn.Module):
+    def __init__(self, c_ch, h_ch):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(c_ch, 2*h_ch),
+            nn.SiLU(),
+            nn.Linear(2*h_ch, 2*h_ch)
+        )
+        # init to identity: gamma≈0 (so 1+gamma=1), beta=0
+        nn.init.zeros_(self.net[-1].weight)
+        nn.init.zeros_(self.net[-1].bias)
+        self.norm = nn.LayerNorm(h_ch)
+
+    def forward(self, h, c):
+        gb = self.net(c)
+        gamma, beta = gb.chunk(2, dim=-1)
+        return (1 + gamma) * self.norm(h) + beta
+
+
+class ConditionalGNOBlock(MessagePassing):
+    """
+    Message passing with FiLM applied to the MESSAGE m_ij,
+    using edge context c_ij = (c_i + c_j)/2.
+    """
+    def __init__(self, hidden_ch, edge_ch=0, aggr="mean"):
+        super().__init__(aggr=aggr, node_dim=0)
+        self.pre_norm = nn.LayerNorm(hidden_ch)
+
+        # raw message builder
+        self.msg = nn.Sequential(
+            nn.Linear(2*hidden_ch + edge_ch, 2*hidden_ch),
+            nn.SiLU(),
+            nn.Linear(2*hidden_ch, hidden_ch)
+        )
+
+        # FiLM over the message (per-edge)
+        self.film_msg = FiLM(c_ch=hidden_ch, h_ch=hidden_ch)
+
+        # node update with residual
+        self.update_mlp = nn.Sequential(
+            nn.Linear(2*hidden_ch, hidden_ch),
+            nn.SiLU(),
+            nn.Linear(hidden_ch, hidden_ch)
+        )
+
+    def forward(self, x, c, edge_index, edge_attr=None):
+        # pre-norm helps stability
+        x_in = x
+        x = self.pre_norm(x)
+        m = self.propagate(edge_index, x=x, c=c, edge_attr=edge_attr)
+        out = self.update_mlp(torch.cat([x_in, m], dim=-1))
+        return x_in + out  # residual
+
+    def message(self, x_i, x_j, c_i, c_j, edge_attr):
+        if edge_attr is not None:
+            m_in = torch.cat([x_i, x_j, edge_attr], dim=-1)
+        else:
+            m_in = torch.cat([x_i, x_j], dim=-1)
+
+        m_raw = self.msg(m_in)
+
+        # edge conditioning: simple mean
+        c_ctx = 0.5 * (c_i + c_j)
+        m = self.film_msg(m_raw, c_ctx)
+        return m
+
+
+class GatingGNO(nn.Module):
+    """
+    In:
+      x : [N, Cx]  (e.g., u or features to predict from)
+      c : [N, Cf]  (conditioning field, e.g., conductivity)
+    Out:
+      y : [N, out_ch]
+    """
+    def __init__(self, x_ch_node, f_ch_node, hidden, layers, edge_ch=0, out_ch=1):
+        super().__init__()
+        self.encoder_x = nn.Sequential(
+            nn.Linear(x_ch_node, hidden // 2),
+            nn.SiLU(),
+            nn.Linear(hidden // 2, hidden),
+        )
+        self.encoder_c = nn.Sequential(
+            nn.Linear(f_ch_node, hidden // 2),
+            nn.SiLU(),
+            nn.Linear(hidden // 2, hidden),
+        )
+        self.blocks = nn.ModuleList(
+            [ConditionalGNOBlock(hidden_ch=hidden, edge_ch=edge_ch) for _ in range(layers)]
+        )
+        self.dec = nn.Sequential(
+            nn.LayerNorm(hidden),
+            nn.SiLU(),
+            nn.Linear(hidden, out_ch)
+        )
+
+    def forward(self, x, c, edge_index, edge_attr=None):
+        x = self.encoder_x(x)  # [N,H]
+        c = self.encoder_c(c)  # [N,H]
+        for blk in self.blocks:
+            x = blk(x, c, edge_index, edge_attr=edge_attr)
+        return self.dec(x)