try a new model

ThermalSolver/model/learnable_finite_difference.py

@@ -2,40 +2,6 @@ import torch
 import torch.nn as nn
 from torch_geometric.nn import MessagePassing
 from torch.nn.utils import spectral_norm
-from matplotlib.tri import Triangulation
-from matplotlib import pyplot as plt
-
-
-def _plot_mesh(y_pred, batch, iteration=None):
-
-    idx = batch.batch == 0
-    y = batch.y[idx].detach().cpu()
-    y_pred = y_pred[idx].detach().cpu()
-    pos = batch.pos[idx].detach().cpu()
-
-    pos = pos.detach().cpu()
-    tria = Triangulation(pos[:, 0], pos[:, 1])
-    plt.figure(figsize=(18, 5))
-    plt.subplot(1, 3, 1)
-    plt.tricontourf(tria, y.squeeze().numpy(), levels=14)
-    plt.colorbar()
-    plt.title("True temperature")
-    plt.subplot(1, 3, 2)
-    plt.tricontourf(tria, y_pred.squeeze().numpy(), levels=14)
-    plt.colorbar()
-    plt.title("Predicted temperature")
-    plt.subplot(1, 3, 3)
-    plt.tricontourf(tria, torch.abs(y_pred - y).squeeze().numpy(), levels=14)
-    plt.colorbar()
-    plt.title("Error")
-    plt.suptitle("GNO", fontsize=16)
-    name = (
-        f"images/gno_iter_{iteration:04d}.png"
-        if iteration is not None
-        else "gno.png"
-    )
-    plt.savefig(name, dpi=72)
-    plt.close()
 
 
 class FiniteDifferenceStep(MessagePassing):
@@ -51,28 +17,12 @@ class FiniteDifferenceStep(MessagePassing):
             spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)),
         )
 
-        self.edge_embedding = nn.Sequential(
-            spectral_norm(nn.Linear(edge_ch, hidden_dim // 2)),
-            nn.GELU(),
-            spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)),
-        )
-
         self.update_net = nn.Sequential(
             spectral_norm(nn.Linear(2 * hidden_dim, hidden_dim)),
             nn.GELU(),
             spectral_norm(nn.Linear(hidden_dim, hidden_dim)),
-            nn.GELU(),
-            # spectral_norm(nn.Linear(hidden_dim // 2, 1)),
         )
 
-        # self.message_net = nn.Sequential(
-        #     spectral_norm(nn.Linear(2 * hidden_dim, hidden_dim)),
-        #     nn.GELU(),
-        #     spectral_norm(nn.Linear(hidden_dim, hidden_dim // 2)),
-        #     nn.GELU(),
-        #     spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)),
-        # )
-
         self.out_net = nn.Sequential(
             spectral_norm(nn.Linear(hidden_dim, hidden_dim // 2)),
             nn.GELU(),
@@ -84,17 +34,14 @@ class FiniteDifferenceStep(MessagePassing):
         TODO: add docstring.
         """
         x_ = self.x_embedding(x)
-        edge_attr_ = self.edge_embedding(edge_attr)
-        out = self.propagate(edge_index, x=x_, edge_attr=edge_attr_, deg=deg)
+        out = self.propagate(edge_index, x=x_, edge_attr=edge_attr, deg=deg)
         return self.out_net(x_ + out)
 
-    def message(self, x_j, edge_attr):
+    def message(self, x_i, x_j, edge_attr):
         """
         TODO: add docstring.
         """
-        # msg_input = torch.cat([x_j, edge_attr], dim=-1)
-        # return self.message_net(msg_input) * edge_attr[:, 3].view(-1, 1)
-        return x_j * edge_attr
+        return (x_j - x_i) * edge_attr.view(-1, 1)
 
     def update(self, aggr_out, x):
         """
@@ -102,10 +49,6 @@ class FiniteDifferenceStep(MessagePassing):
         """
         update_input = torch.cat([x, aggr_out], dim=-1)
         return self.update_net(update_input)
-        # return self.update_net(aggr_out)
-        # return aggr_out
-        # h = self.update_net(aggr_out, x)
-        # return h
 
     def aggregate(self, inputs, index, deg):
         """
@@ -114,13 +57,3 @@ class FiniteDifferenceStep(MessagePassing):
         out = super().aggregate(inputs, index)
         deg = deg + 1e-7
         return out / deg.view(-1, 1)
-
-
-# # Da fare:
-# # - Finire calcolo della loss su ogni step e poi media
-# # - Test con vari modelli
-# # - Se non dovesse funzionare, provare ad adeguare il criterio di uscita
-
-# # PINN batching:
-# # - Provare singola condizione
-# # - Ottimizzatore del secondo ordine (LBFGS)