implement ML correction

ThermalSolver/model/learnable_finite_difference.py

@@ -1,53 +1,53 @@
 import torch
 import torch.nn as nn
 from torch_geometric.nn import MessagePassing
-from torch.nn.utils import spectral_norm
+
+# from torch.nn.utils import spectral_norm
 
 
-class FiniteDifferenceStep(MessagePassing):
-    """
-    TODO: add docstring.
-    """
-
-    def __init__(self, hidden_dim=16, aggr: str = "add"):
-        print(aggr)
-        super().__init__(aggr=aggr)
-        self.x_embedding = nn.Sequential(
-            spectral_norm(nn.Linear(1, hidden_dim // 2)),
-            nn.GELU(),
-            spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)),
+class GCNConvLayer(MessagePassing):
+    def __init__(self, in_channels, out_channels):
+        super().__init__("add")
+        self.lin = nn.Sequential(
+            nn.Linear(in_channels, out_channels),
+            nn.ReLU(),
+            nn.Linear(out_channels, out_channels),
+            nn.ReLU(),
         )
 
-        self.out_net = nn.Sequential(
-            spectral_norm(nn.Linear(hidden_dim, hidden_dim // 2)),
-            nn.GELU(),
-            spectral_norm(nn.Linear(hidden_dim // 2, 1)),
+    def _compute_edge_weight(self, edge_index, edge_w, deg):
+        """ """
+        return edge_w.squeeze() / (
+            1 + torch.sqrt(deg[edge_index[0]] * deg[edge_index[1]])
         )
 
     def forward(self, x, edge_index, edge_attr, deg):
-        """
-        TODO: add docstring.
-        """
-        x_ = self.x_embedding(x)
-        out = self.propagate(edge_index, x=x_, edge_attr=edge_attr, deg=deg)
-        return self.out_net(out)
+        edge_w = self._compute_edge_weight(edge_index, edge_attr, deg)
+        return self.propagate(edge_index, x=x, edge_weight=edge_w, deg=deg)
 
-    def message(self, x_j, edge_attr):
-        """
-        TODO: add docstring.
-        """
-        return x_j * edge_attr.view(-1, 1)
+    def message(self, x_j, edge_weight):
+        return edge_weight.view(-1, 1) * x_j
 
-    def update(self, aggr_out, _):
-        """
-        TODO: add docstring.
-        """
-        return aggr_out
 
-    def aggregate(self, inputs, index, deg):
-        """
-        TODO: add docstring.
-        """
-        out = super().aggregate(inputs, index)
-        deg = deg + 1e-7
-        return out / deg.view(-1, 1)
+class CorrectionNet(nn.Module):
+    def __init__(self, hidden_dim=8):
+        super().__init__()
+        self.enc = nn.Sequential(
+            nn.Linear(1, hidden_dim // 2),
+            nn.ReLU(),
+            nn.Linear(hidden_dim // 2, hidden_dim),
+            nn.ReLU(),
+        )
+        self.model = GCNConvLayer(hidden_dim, hidden_dim)
+        self.dec = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.ReLU(),
+            nn.Linear(hidden_dim // 2, 1),
+            nn.ReLU(),
+        )
+
+    def forward(self, x, edge_index, edge_attr, deg):
+        h = self.enc(x)
+        h = self.model(h, edge_index, edge_attr, deg)
+        out = self.dec(h)
+        return out