add config file

ThermalSolver/graph_datamodule.py

@@ -18,6 +18,8 @@ class GraphDataModule(LightningDataModule):
         test_size: float = 0.1,
         batch_size: int = 32,
         remove_boundary_edges: bool = False,
+        build_radial_graph: bool = False,
+        radius: float = None,
     ):
         super().__init__()
         self.hf_repo = hf_repo
@@ -29,6 +31,8 @@ class GraphDataModule(LightningDataModule):
         self.test_size = test_size
         self.batch_size = batch_size
         self.remove_boundary_edges = remove_boundary_edges
+        self.build_radial_graph = build_radial_graph
+        self.radius = radius
 
     def prepare_data(self):
         dataset = load_dataset(self.hf_repo, name="snapshots")[self.split_name]
@@ -80,9 +84,8 @@ class GraphDataModule(LightningDataModule):
         )
         temperature = torch.tensor(snapshot["temperature"], dtype=torch.float32)
 
-        edge_index = torch.tensor(geometry["edge_index"], dtype=torch.int64).T
-
         pos = torch.tensor(geometry["points"], dtype=torch.float32)[:, :2]
+
         bottom_ids = torch.tensor(
             geometry["bottom_boundary_ids"], dtype=torch.long
         )
@@ -92,20 +95,38 @@ class GraphDataModule(LightningDataModule):
             geometry["right_boundary_ids"], dtype=torch.long
         )
 
-        edge_index = to_undirected(edge_index, num_nodes=pos.size(0))
+        if self.build_radial_graph:
+            from pina.graph import RadiusGraph
+
+            if self.radius is None:
+                raise ValueError("Radius must be specified for radial graph.")
+            edge_index = RadiusGraph.compute_radius_graph(
+                pos, radius=self.radius
+            )
+            from torch_geometric.utils import remove_self_loops
+
+            edge_index, _ = remove_self_loops(edge_index)
+        else:
+
+            edge_index = torch.tensor(
+                geometry["edge_index"], dtype=torch.int64
+            ).T
+            edge_index = to_undirected(edge_index, num_nodes=pos.size(0))
 
         boundary_mask, boundary_values = self._compute_boundary_mask(
             bottom_ids, right_ids, top_ids, left_ids, temperature
         )
+
         if self.remove_boundary_edges:
             boundary_idx = torch.unique(boundary_mask)
             edge_index_mask = ~torch.isin(edge_index[1], boundary_idx)
             edge_index = edge_index[:, edge_index_mask]
 
-        edge_attr = pos[edge_index[0]] - pos[edge_index[1]]
+        # edge_attr = pos[edge_index[0]] - pos[edge_index[1]]
-        edge_attr = torch.cat(
+        # edge_attr = torch.cat(
-            [edge_attr, torch.norm(edge_attr, dim=1).unsqueeze(-1)], dim=1
+        #     [edge_attr, torch.norm(edge_attr, dim=1).unsqueeze(-1)], dim=1
-        )
+        # )
+        edge_attr = torch.norm(pos[edge_index[0]] - pos[edge_index[1]], dim=1)
 
         x = torch.zeros_like(temperature, dtype=torch.float32).unsqueeze(-1)
         if self.remove_boundary_edges:

ThermalSolver/graph_module.py

@@ -4,6 +4,7 @@ from torch_geometric.data import Batch
 import importlib
 from matplotlib import pyplot as plt
 from matplotlib.tri import Triangulation
+from .model.finite_difference import FiniteDifferenceStep
 
 
 def import_class(class_path: str):
@@ -13,7 +14,7 @@ def import_class(class_path: str):
     return cls
 
 
-def _plot_mesh(pos, y, y_pred, batch):
+def _plot_mesh(pos, y, y_pred, batch, i):
 
     idx = batch == 0
     y = y[idx].detach().cpu()
@@ -36,48 +37,48 @@ def _plot_mesh(pos, y, y_pred, batch):
     plt.colorbar()
     plt.title("Error")
     plt.suptitle("GNO", fontsize=16)
-    plt.savefig("gno.png", dpi=300)
+    name = f"images/graph_iter_{i:04d}.png"
+    plt.savefig(name, dpi=72)
+    plt.close()
 
 
 class GraphSolver(LightningModule):
     def __init__(
         self,
         model_class_path: str,
-        model_init_args: dict,
+        model_init_args: dict = {},
         loss: torch.nn.Module = None,
-        unrolling_steps: int = 48,
+        curriculum_learning: bool = False,
+        start_iters: int = 10,
+        increase_every: int = 100,
+        increase_rate: float = 1.1,
+        max_iters: int = 1000,
+        accumulation_iters: int = None,
     ):
         super().__init__()
         self.model = import_class(model_class_path)(**model_init_args)
+        self.fd_net = FiniteDifferenceStep()
         self.loss = loss if loss is not None else torch.nn.MSELoss()
-        self.unrolling_steps = unrolling_steps
+        self.curriculum_learning = curriculum_learning
+        self.start_iters = start_iters
+        self.increase_every = increase_every
+        self.increase_rate = increase_rate
+        self.max_iters = max_iters
+        self.current_iters = start_iters
+        self.accumulation_iters = accumulation_iters
+        self.automatic_optimization = False
+        self.threshold = 1e-5
 
-    def forward(
+        self.alpha = torch.nn.Parameter(torch.tensor(0.1))
-        self,
+
-        x: torch.Tensor,
+    def _compute_deg(self, edge_index, edge_attr, num_nodes):
-        c: torch.Tensor,
+        deg = torch.zeros(num_nodes, device=edge_index.device)
-        edge_index: torch.Tensor,
+        deg = deg.scatter_add(0, edge_index[1], edge_attr)
-        edge_attr: torch.Tensor,
+        return deg + 1e-7
-        unrolling_steps: int = None,
-        boundary_mask: torch.Tensor = None,
-        boundary_values: torch.Tensor = None,
-    ):
-        return self.model(
-            x=x,
-            c=c,
-            edge_index=edge_index,
-            edge_attr=edge_attr,
-            unrolling_steps=unrolling_steps,
-            boundary_mask=boundary_mask,
-            boundary_values=boundary_values,
-        )
 
     def _compute_loss(self, x, y):
         return self.loss(x, y)
 
-    def _preprocess_batch(self, batch: Batch):
-        return batch.x, batch.y, batch.c, batch.edge_index, batch.edge_attr
-
     def _log_loss(self, loss, batch, stage: str):
         self.log(
             f"{stage}/loss",
@@ -89,89 +90,206 @@ class GraphSolver(LightningModule):
         )
         return loss
 
+    @staticmethod
+    def _compute_c_ij(c, edge_index):
+        """
+        TODO: add docstring.
+        """
+        return (0.5 * (c[edge_index[0]] + c[edge_index[1]])).squeeze()
+
+    def _compute_model_steps(
+        self, x, edge_index, edge_attr, deg, boundary_mask, boundary_values
+    ):
+        # with torch.no_grad():
+        # out = self.fd_net(x, edge_index, edge_attr, deg)
+        # out[boundary_mask] = boundary_values.unsqueeze(-1)
+        # diff = out - x
+        # out = self.model(out, edge_index, edge_attr, deg)
+        # out = out + self.alpha * correction
+        # out[boundary_mask] = boundary_values.unsqueeze(-1)
+        out = self.model(x, edge_index, edge_attr, deg)
+        out[boundary_mask] = boundary_values.unsqueeze(-1)
+        return out
+
+    def _check_convergence(self, out, x):
+        residual_norm = torch.norm(out - x)
+        if residual_norm < self.threshold * torch.norm(x):
+            return True
+        return False
+
+    def accumulate_gradients(self, losses):
+        loss_ = torch.stack(losses, dim=0).mean()
+        self.manual_backward(loss_, retain_graph=True)
+        return loss_.item()
+
+    def _preprocess_batch(self, batch: Batch):
+        x, y, c, edge_index, edge_attr = (
+            batch.x,
+            batch.y,
+            batch.c,
+            batch.edge_index,
+            batch.edge_attr,
+        )
+        edge_attr = 1 / edge_attr
+        c_ij = self._compute_c_ij(c, edge_index)
+        edge_attr = edge_attr * c_ij
+        return x, y, edge_index, edge_attr
+
     def training_step(self, batch: Batch, _):
-        x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
+        optim = self.optimizers()
-        y_pred, it = self(
+        optim.zero_grad()
-            x,
+        x, y, edge_index, edge_attr = self._preprocess_batch(batch)
-            c,
+
-            edge_index=edge_index,
+        deg = self._compute_deg(edge_index, edge_attr, x.size(0))
-            edge_attr=edge_attr,
+        losses = []
-            unrolling_steps=self.unrolling_steps,
+        acc_loss, acc_it = 0, 0
-            boundary_mask=batch.boundary_mask,
+
-            boundary_values=batch.boundary_values,
+        for i in range(self.current_iters):
+            out = self._compute_model_steps(
+                x,
+                edge_index,
+                edge_attr.unsqueeze(-1),
+                deg,
+                batch.boundary_mask,
+                batch.boundary_values,
+            )
+            losses.append(self.loss(out, y))
+
+            # Accumulate gradients if reached accumulation iters
+            if (
+                self.accumulation_iters is not None
+                and (i + 1) % self.accumulation_iters == 0
+            ):
+                loss = self.accumulate_gradients(losses)
+                losses = []
+                acc_it += 1
+                out = out.detach()
+                acc_loss = acc_loss + loss
+
+            # Check for convergence and break if converged (with final accumulation)
+            converged = self._check_convergence(out, x)
+            if converged:
+                if losses:
+                    loss = self.accumulate_gradients(losses)
+                    acc_it += 1
+                    acc_loss = acc_loss + loss
+                break
+
+            # Final accumulation if we are at the last iteration
+            if i == self.current_iters - 1:
+                if losses:
+                    loss = self.accumulate_gradients(losses)
+                    acc_it += 1
+                    acc_loss = acc_loss + loss
+
+            x = out
+
+        loss = self.loss(out, y)
+        for param in self.model.parameters():
+            if param.grad is not None:
+                param.grad /= acc_it
+        optim.step()
+        optim.zero_grad()
+
+        self.log(
+            "train/accumulated_loss",
+            (acc_loss / acc_it if acc_it > 0 else acc_loss),
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            batch_size=int(batch.num_graphs),
         )
-        loss = self.loss(y_pred, y)
-        boundary_loss = self.loss(
-            y_pred[batch.boundary_mask], y[batch.boundary_mask]
-        )
-        self._log_loss(loss, batch, "train")
-        # self._log_loss(boundary_loss, batch, "train_boundary")
         self.log(
             "train/iterations",
-            it,
+            i + 1,
             on_step=False,
             on_epoch=True,
             prog_bar=True,
             batch_size=int(batch.num_graphs),
         )
-        self.log(
+        if hasattr(self.model, "p"):
-            "train/param_p",
+            self.log(
-            self.model.fd_step.p,
+                "train/p",
-            on_step=False,
+                self.model.p,
-            on_epoch=True,
+                on_step=False,
-            prog_bar=True,
+                on_epoch=True,
-            batch_size=int(batch.num_graphs),
+                prog_bar=True,
-        )
+                batch_size=int(batch.num_graphs),
-        # self.log("train/param_a", self.model.fd_step.a, on_step=False, on_epoch=True, prog_bar=True, batch_size=int(batch.num_graphs))
+            )
-        return loss
+
+    def on_train_epoch_end(self):
+        if self.curriculum_learning:
+            if (self.current_iters < self.max_iters) and (
+                self.current_epoch % self.increase_every == 0
+            ):
+                self.current_iters = min(
+                    int(self.current_iters * self.increase_rate), self.max_iters
+                )
+        return super().on_train_epoch_end()
 
     def validation_step(self, batch: Batch, _):
-        x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
+        x, y, edge_index, edge_attr = self._preprocess_batch(batch)
-        y_pred, it = self(
+
-            x,
+        deg = self._compute_deg(edge_index, edge_attr, x.size(0))
-            c,
+
-            edge_index=edge_index,
+        for i in range(self.current_iters):
-            edge_attr=edge_attr,
+            out = self._compute_model_steps(
-            unrolling_steps=self.unrolling_steps,
+                x,
-            boundary_mask=batch.boundary_mask,
+                edge_index,
-            boundary_values=batch.boundary_values,
+                edge_attr.unsqueeze(-1),
-        )
+                deg,
-        loss = self.loss(y_pred, y)
+                batch.boundary_mask,
-        boundary_loss = self.loss(
+                batch.boundary_values,
-            y_pred[batch.boundary_mask], y[batch.boundary_mask]
+            )
-        )
+            converged = self._check_convergence(out, x)
+
+            if converged:
+                break
+            x = out
+        loss = self.loss(out, y)
         self._log_loss(loss, batch, "val")
         self.log(
             "val/iterations",
-            it,
+            i + 1,
             on_step=False,
             on_epoch=True,
             prog_bar=True,
             batch_size=int(batch.num_graphs),
         )
-        return loss
 
     def test_step(self, batch: Batch, _):
-        x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
+        x, y, edge_index, edge_attr = self._preprocess_batch(batch)
-        y_pred, _ = self.model(
+
-            x=x,
+        deg = self._compute_deg(edge_index, edge_attr, x.size(0))
-            c=c,
+
-            edge_index=edge_index,
+        for i in range(self.max_iters):
-            edge_attr=edge_attr,
+            out = self._compute_model_steps(
-            unrolling_steps=self.unrolling_steps,
+                x,
-            batch=batch.batch,
+                edge_index,
-            pos=batch.pos,
+                edge_attr.unsqueeze(-1),
-            boundary_mask=batch.boundary_mask,
+                deg,
-            boundary_values=batch.boundary_values,
+                batch.boundary_mask,
-            plot_results=False,
+                batch.boundary_values,
-        )
+            )
-        loss = self._compute_loss(y_pred, y)
+            converged = self._check_convergence(out, x)
-        _plot_mesh(batch.pos, y, y_pred, batch.batch)
+            # _plot_mesh(batch.pos, y, out, batch.batch, i)
+            if converged:
+                break
+            x = out
+        loss = self.loss(out, y)
+
         self._log_loss(loss, batch, "test")
-        return loss
+        self.log(
+            "test/iterations",
+            i + 1,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            batch_size=int(batch.num_graphs),
+        )
 
     def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        optimizer = torch.optim.AdamW(self.parameters(), lr=1e-3)
         return optimizer
 
     def _impose_bc(self, x: torch.Tensor, data: Batch):

ThermalSolver/model/__init__.py

@@ -1,13 +1,13 @@
 __all__ = [
-    "GraphFiniteDifference",
+    # "GraphFiniteDifference",
     "GatingGNO",
-    "LearnableGraphFiniteDifference",
+    # "LearnableGraphFiniteDifference",
     "PointNet",
 ]
 
-from .learnable_finite_difference import (
+# from .learnable_finite_difference import (
-    GraphFiniteDifference as LearnableGraphFiniteDifference,
+# GraphFiniteDifference as LearnableGraphFiniteDifference,
-)
+# )
-from .finite_difference import GraphFiniteDifference as GraphFiniteDifference
+# from .finite_difference import GraphFiniteDifference as GraphFiniteDifference
 from .local_gno import GatingGNO
 from .point_net import PointNet

ThermalSolver/model/finite_difference.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 from torch_geometric.nn import MessagePassing
+from torch.nn.utils import spectral_norm
 
 
 class FiniteDifferenceStep(MessagePassing):
@@ -8,14 +9,8 @@ class FiniteDifferenceStep(MessagePassing):
     TODO: add docstring.
     """
 
-    def __init__(self, aggr: str = "add", root_weight: float = 1.0):
+    def __init__(self):
-        super().__init__(aggr=aggr)
+        super().__init__(aggr="add")
-        assert (
-            aggr == "add"
-        ), "Per somme pesate, l'aggregazione deve essere 'add'."
-        # self.root_weight = float(root_weight)
-        self.p = torch.nn.Parameter(torch.tensor(0.8))
-        self.a = root_weight
 
     def forward(self, x, edge_index, edge_attr, deg):
         """
@@ -28,8 +23,13 @@ class FiniteDifferenceStep(MessagePassing):
         """
         TODO: add docstring.
         """
-        p = torch.clamp(self.p, 0.0, 1.0)
+        return x_j * edge_attr
-        return p * edge_attr.view(-1, 1) * x_j
+
+    def update(self, aggr_out, _):
+        """
+        TODO: add docstring.
+        """
+        return aggr_out
 
     def aggregate(self, inputs, index, deg):
         """
@@ -38,84 +38,3 @@ class FiniteDifferenceStep(MessagePassing):
         out = super().aggregate(inputs, index)
         deg = deg + 1e-7
         return out / deg.view(-1, 1)
-
-    def update(self, aggr_out, x):
-        """
-        TODO: add docstring.
-        """
-        a = torch.clamp(self.a, 0.0, 1.0)
-        return a * aggr_out + (1 - a) * x
-        # return self.a * aggr_out + (1 - self.a) * x
-
-
-class GraphFiniteDifference(nn.Module):
-    """
-    TODO: add docstring.
-    """
-
-    def __init__(self, max_iters: int = 5000, threshold: float = 1e-4):
-        """
-        TODO: add docstring.
-        """
-        super().__init__()
-        self.max_iters = max_iters
-        self.threshold = threshold
-        self.fd_step = FiniteDifferenceStep(aggr="add", root_weight=1.0)
-
-    @staticmethod
-    def _compute_deg(edge_index, edge_attr, num_nodes):
-        """
-        TODO: add docstring.
-        """
-        deg = torch.zeros(num_nodes, device=edge_index.device)
-        deg = deg.scatter_add(0, edge_index[1], edge_attr)
-        return deg + 1e-7
-
-    @staticmethod
-    def _compute_c_ij(c, edge_index):
-        """
-        TODO: add docstring.
-        """
-        return (0.5 * (c[edge_index[0]] + c[edge_index[1]])).squeeze()
-
-    def forward(
-        self,
-        x,
-        edge_index,
-        edge_attr,
-        c,
-        boundary_mask,
-        boundary_values,
-        **kwargs,
-    ):
-        """
-        TODO: add docstring.
-        """
-        edge_attr = 1 / edge_attr[:, -1]
-        c_ij = self._compute_c_ij(c, edge_index)
-        edge_attr = edge_attr * c_ij
-        deg = self._compute_deg(edge_index, edge_attr, x.size(0))
-
-        # Calcola la soglia staccando x dal grafo
-        conv_thres = self.threshold * torch.norm(x.detach())
-
-        for _i in range(self.max_iters):
-            out = self.fd_step(x, edge_index, edge_attr, deg)
-            out[boundary_mask] = boundary_values.unsqueeze(-1)
-
-            # Controllo convergenza senza tracciamento gradienti
-            with torch.no_grad():
-                residual_norm = torch.norm(out - x)
-
-            if residual_norm < conv_thres:
-                break
-
-            # --- OTTIMIZZAZIONE CHIAVE ---
-            # Stacca 'out' dal grafo prima della prossima iterazione
-            # per evitare BPTT e risparmiare memoria.
-            x = out.detach()
-
-        # Il 'out' finale restituito mantiene i gradienti
-        # dell'ULTIMA chiamata a fd_step, permettendo al modello
-        # di apprendere correttamente.
-        return out, _i + 1

ThermalSolver/model/learnable_finite_difference.py

@@ -1,58 +1,119 @@
+# import torch
+# import torch.nn as nn
+# from torch_geometric.nn import MessagePassing
+# from torch.nn.utils import spectral_norm
+
+# class GCNConvLayer(MessagePassing):
+#     def __init__(self, in_channels, out_channels):
+#         super().__init__(aggr="add")
+#         self.lin_l = spectral_norm(nn.Linear(in_channels, out_channels, bias=False))
+#         self.lin_r = spectral_norm(nn.Linear(in_channels, out_channels, bias=False))
+
+#     def forward(self, x, edge_index, edge_attr, deg):
+#         out = self.propagate(edge_index, x=x, edge_attr=edge_attr, deg=deg)
+#         out = self.lin_l(out)
+#         return out
+
+#     def message(self, x_j, edge_attr):
+#         return x_j * edge_attr
+
+#     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, n_layers=1):
+#         super().__init__()
+#         # self.enc = GCNConvLayer(1, hidden_dim)
+#         self.enc = nn.Sequential(
+#             spectral_norm(nn.Linear(1, hidden_dim//2)),
+#             nn.GELU(),
+#             spectral_norm(nn.Linear(hidden_dim//2, hidden_dim)),
+#         )
+#         self.layers = torch.nn.ModuleList([GCNConvLayer(hidden_dim, hidden_dim) for _ in range(n_layers)])
+#         self.relu = nn.GELU()
+
+#         self.dec = nn.Sequential(
+#             spectral_norm(nn.Linear(hidden_dim, hidden_dim//2)),
+#             nn.GELU(),
+#             spectral_norm(nn.Linear(hidden_dim//2, 1)),
+#         )
+
+#     def forward(self, x, edge_index, edge_attr, deg,):
+#         # h = self.enc(x, edge_index, edge_attr, deg)
+#         # h = self.relu(self.enc(x))
+#         h = self.enc(x)
+#         for layer in self.layers:
+#             h = layer(h, edge_index, edge_attr, deg)
+#             # h = self.norm(h)
+#             h = self.relu(h)
+#         # out = self.dec(h, edge_index, edge_attr, deg)
+#         out = self.dec(h)
+#         return out
+
+
 import torch
 import torch.nn as nn
 from torch_geometric.nn import MessagePassing
 from torch.nn.utils import spectral_norm
 
 
-class FiniteDifferenceStep(MessagePassing):
+class CorrectionNet(MessagePassing):
     """
     TODO: add docstring.
     """
 
-    def __init__(self, aggr: str = "add", root_weight: float = 1.0):
+    def __init__(self, hidden_dim=16):
-        super().__init__(aggr=aggr)
+        super().__init__(aggr="add")
-        assert (
+        self.in_net = nn.Sequential(
-            aggr == "add"
+            spectral_norm(nn.Linear(1, hidden_dim // 2)),
-        ), "Per somme pesate, l'aggregazione deve essere 'add'."
+            nn.GELU(),
-
+            spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)),
-        self.correction_net = nn.Sequential(
-            nn.Linear(2, 6),
-            nn.Tanh(),
-            nn.Linear(6, 1),
-            nn.Tanh(),
-        )
-        self.update_net = nn.Sequential(
-            spectral_norm(nn.Linear(1, 6)),
-            nn.Softplus(),
-            spectral_norm(nn.Linear(6, 1)),
-            nn.Softplus(),
         )
 
-        self.message_net = nn.Sequential(
+        self.out_net = nn.Sequential(
-            spectral_norm(nn.Linear(1, 6)),
+            spectral_norm(nn.Linear(hidden_dim, hidden_dim // 2)),
-            nn.Softplus(),
+            nn.GELU(),
-            spectral_norm(nn.Linear(6, 1)),
+            spectral_norm(nn.Linear(hidden_dim // 2, 1)),
-            nn.Softplus(),
         )
-        self.p = torch.nn.Parameter(torch.tensor(0.5))
+
-        # self.a = torch.nn.Parameter(torch.tensor(root_weight))
+        self.lin_msg = spectral_norm(
+            nn.Linear(hidden_dim, hidden_dim, bias=False)
+        )
+        self.lin_update = spectral_norm(
+            nn.Linear(hidden_dim, hidden_dim, bias=False)
+        )
+        self.alpha = nn.Parameter(torch.tensor(0.0))
+        self.beta = nn.Parameter(torch.tensor(0.0))
 
     def forward(self, x, edge_index, edge_attr, deg):
         """
         TODO: add docstring.
         """
+        x = self.in_net(x)
         out = self.propagate(edge_index, x=x, edge_attr=edge_attr, deg=deg)
-        return out
+        return self.out_net(out)
 
     def message(self, x_j, edge_attr):
         """
         TODO: add docstring.
         """
-        # x_in = torch.cat([x_j, edge_attr.view(-1, 1)], dim=-1)
+        alpha = torch.sigmoid(self.alpha)
-        # correction = self.correction_net(x_in)
+        msg = x_j * edge_attr
-        # p = torch.sigmoid(self.p)
+        msg = (1 - alpha) * msg + alpha * self.lin_msg(msg)
-        # return (p * edge_attr.view(-1, 1) + (1 - p) * correction) * x_j
+        return msg
-        return edge_attr.view(-1, 1) * x_j
+
+    def update(self, aggr_out, x):
+        """
+        TODO: add docstring.
+        """
+        beta = torch.sigmoid(self.beta)
+        return aggr_out * (1 - beta) + self.lin_msg(x) * beta
 
     def aggregate(self, inputs, index, deg):
         """
@@ -61,69 +122,3 @@ class FiniteDifferenceStep(MessagePassing):
         out = super().aggregate(inputs, index)
         deg = deg + 1e-7
         return out / deg.view(-1, 1)
-
-    def update(self, aggr_out, x):
-        """
-        TODO: add docstring.
-        """
-        return self.update_net(aggr_out)
-
-
-class GraphFiniteDifference(nn.Module):
-    """
-    TODO: add docstring.
-    """
-
-    def __init__(self, max_iters: int = 5000, threshold: float = 1e-4):
-        """
-        TODO: add docstring.
-        """
-        super().__init__()
-        self.max_iters = max_iters
-        self.threshold = threshold
-        self.fd_step = FiniteDifferenceStep(aggr="add", root_weight=1.0)
-
-    @staticmethod
-    def _compute_deg(edge_index, edge_attr, num_nodes):
-        """
-        TODO: add docstring.
-        """
-        deg = torch.zeros(num_nodes, device=edge_index.device)
-        deg = deg.scatter_add(0, edge_index[1], edge_attr)
-        return deg + 1e-7
-
-    @staticmethod
-    def _compute_c_ij(c, edge_index):
-        """
-        TODO: add docstring.
-        """
-        return (0.5 * (c[edge_index[0]] + c[edge_index[1]])).squeeze()
-
-    def forward(
-        self,
-        x,
-        edge_index,
-        edge_attr,
-        c,
-        boundary_mask,
-        boundary_values,
-        **kwargs,
-    ):
-        """
-        TODO: add docstring.
-        """
-        edge_attr = 1 / edge_attr[:, -1]
-        c_ij = self._compute_c_ij(c, edge_index)
-        edge_attr = edge_attr * c_ij
-        deg = self._compute_deg(edge_index, edge_attr, x.size(0))
-        conv_thres = self.threshold * torch.norm(x.detach())
-
-        for _i in range(self.max_iters):
-            out = self.fd_step(x, edge_index, edge_attr, deg)
-            out[boundary_mask] = boundary_values.unsqueeze(-1)
-            with torch.no_grad():
-                residual_norm = torch.norm(out - x)
-            if residual_norm < conv_thres:
-                break
-            x = out.detach()
-        return out, _i + 1

experiments/config_gno.yaml

@@ -0,0 +1,62 @@
+# lightning.pytorch==2.5.5
+seed_everything: 1999
+trainer:
+  accelerator: gpu
+  strategy: auto
+  devices: 1
+  num_nodes: 1
+  precision: null
+  logger:
+    - class_path: lightning.pytorch.loggers.TensorBoardLogger
+      init_args:
+        save_dir: logs
+        name: "test"
+        version: null
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        monitor: val/loss
+        mode: min
+        save_top_k: 1
+        filename: best-checkpoint
+    - class_path: lightning.pytorch.callbacks.EarlyStopping
+      init_args:
+        monitor: val/loss
+        mode: min
+        patience: 25
+        verbose: false
+  max_epochs: 1000
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  overfit_batches: 0.0
+  log_every_n_steps: null
+  # inference_mode: true
+  default_root_dir: null
+  # accumulate_grad_batches: 2
+  # gradient_clip_val: 1.0
+model: 
+  class_path: ThermalSolver.graph_module.GraphSolver
+  init_args:
+    model_class_path: ThermalSolver.model.learnable_finite_difference.CorrectionNet
+    curriculum_learning: true
+    start_iters: 5
+    increase_every: 10
+    increase_rate: 2
+    max_iters: 2000
+    accumulation_iters: 320
+data: 
+  class_path: ThermalSolver.graph_datamodule.GraphDataModule
+  init_args:
+    hf_repo: "SISSAmathLab/thermal-conduction"
+    split_name: "1000_40x30"
+    batch_size: 32
+    train_size: 0.8
+    test_size: 0.1
+    test_size: 0.1
+    build_radial_graph: false
+    radius: 0.6
+    remove_boundary_edges: false
+optimizer: null
+lr_scheduler: null
+# ckpt_path: logs/test/version_0/checkpoints/best-checkpoint.ckpt