improve unrolling

ThermalSolver/data_module.py

@@ -5,6 +5,7 @@ from datasets import load_dataset
 from torch_geometric.data import Data
 from torch_geometric.loader import DataLoader
 from torch_geometric.utils import to_undirected
+from .mesh_data import MeshData
 
 
 class GraphDataModule(LightningDataModule):
@@ -12,7 +13,7 @@ class GraphDataModule(LightningDataModule):
         self,
         hf_repo: str,
         split_name: str,
-        train_size: float = 0.8,
+        train_size: float = 0.2,
         val_size: float = 0.1,
         test_size: float = 0.1,
         batch_size: int = 32,
@@ -40,45 +41,79 @@ class GraphDataModule(LightningDataModule):
         pos = torch.tensor(self.geometry["points"][0], dtype=torch.float32)[
             :, :2
         ]
-        bottom_boundary_ids = torch.tensor(
+
-            self.geometry["bottom_boundary_ids"][0], dtype=torch.int64
+        bottom_ids = torch.tensor(
+            self.geometry["bottom_boundary_ids"][0], dtype=torch.long
+        )
+        top_ids = torch.tensor(
+            self.geometry["top_boundary_ids"][0], dtype=torch.long
+        )
+        left_ids = torch.tensor(
+            self.geometry["left_boundary_ids"][0], dtype=torch.long
+        )
+        right_ids = torch.tensor(
+            self.geometry["right_boundary_ids"][0], dtype=torch.long
         )
         self.data = [
             self._build_dataset(
-                torch.tensor(snapshot["conductivity"], dtype=torch.float32),
+                snapshot,
-                torch.tensor(snapshot["boundary_values"], dtype=torch.float32),
-                torch.tensor(snapshot["temperature"], dtype=torch.float32),
                 edge_index.T,
                 pos,
-                bottom_boundary_ids,
+                bottom_ids,
+                top_ids,
+                left_ids,
+                right_ids,
             )
             for snapshot in tqdm(hf_dataset, desc="Building graphs")
         ]
 
     def _build_dataset(
         self,
-        conductivity: torch.Tensor,
+        snapshot: dict,
-        boundary_vales: torch.Tensor,
-        temperature: torch.Tensor,
         edge_index: torch.Tensor,
         pos: torch.Tensor,
-        bottom_boundary_ids: torch.Tensor,
+        bottom_ids: torch.Tensor,
+        top_ids: torch.Tensor,
+        left_ids: torch.Tensor,
+        right_ids: torch.Tensor,
     ) -> Data:
+        conductivity = torch.tensor(
+            snapshot["conductivity"], dtype=torch.float32
+        )
+        temperature = torch.tensor(snapshot["temperature"], dtype=torch.float32)
+
         edge_index = to_undirected(edge_index, num_nodes=pos.size(0))
         edge_attr = pos[edge_index[0]] - pos[edge_index[1]]
         edge_attr = torch.cat(
             [edge_attr, torch.norm(edge_attr, dim=1).unsqueeze(-1)], dim=1
         )
-        boundary_temperature = boundary_vales[bottom_boundary_ids].max()
+
-        boundary_vales[bottom_boundary_ids] = 1.0
+        left_ids = left_ids[~torch.isin(left_ids, bottom_ids)]
-        return Data(
+        right_ids = right_ids[~torch.isin(right_ids, bottom_ids)]
-            x=boundary_vales.unsqueeze(-1),
+        left_ids = left_ids[~torch.isin(left_ids, top_ids)]
+        right_ids = right_ids[~torch.isin(right_ids, top_ids)]
+
+        bottom_bc = temperature[bottom_ids].median()
+        bottom_bc_mask = torch.ones(len(bottom_ids)) * bottom_bc
+        left_bc = temperature[left_ids].median()
+        left_bc_mask = torch.ones(len(left_ids)) * left_bc
+        right_bc = temperature[right_ids].median()
+        right_bc_mask = torch.ones(len(right_ids)) * right_bc
+
+        boundary_values = torch.cat(
+            [bottom_bc_mask, right_bc_mask, left_bc_mask], dim=0
+        )
+        boundary_mask = torch.cat([bottom_ids, right_ids, left_ids], dim=0)
+
+        return MeshData(
+            x=torch.rand_like(temperature).unsqueeze(-1),
             c=conductivity.unsqueeze(-1),
             edge_index=edge_index,
             pos=pos,
             edge_attr=edge_attr,
+            boundary_mask=boundary_mask,
+            boundary_values=boundary_values.unsqueeze(-1),
             y=temperature.unsqueeze(-1),
-            boundary_temperature=boundary_vales[bottom_boundary_ids].max(),
         )
 
     def setup(self, stage: str = None):
@@ -92,13 +127,18 @@ class GraphDataModule(LightningDataModule):
         if stage == "test" or stage is None:
             self.test_data = self.data[val_end:]
 
-    def train_dataloader(self) -> DataLoader:
+    # nel tuo LightningDataModule
+    def train_dataloader(self):
         return DataLoader(
             self.train_data, batch_size=self.batch_size, shuffle=True
         )
 
-    def val_dataloader(self) -> DataLoader:
+    def val_dataloader(self):
-        return DataLoader(self.val_data, batch_size=self.batch_size)
+        return DataLoader(
+            self.val_data, batch_size=self.batch_size, shuffle=False
+        )
 
-    def test_dataloader(self) -> DataLoader:
+    def test_dataloader(self):
-        return DataLoader(self.test_data, batch_size=self.batch_size)
+        return DataLoader(
+            self.test_data, batch_size=self.batch_size, shuffle=False
+        )

ThermalSolver/mesh_data.py

@@ -0,0 +1,17 @@
+"""
+Custom Data/Batch per gestire bene le boundary conditions.
+"""
+
+from typing import List
+import torch
+from torch_geometric.data import Data, Batch
+
+B_KEYS: List[str] = ["boundary_mask"]
+
+
+class MeshData(Data):
+    def __inc__(self, key, value, *args, **kwargs):
+        # questi campi sono INDICI di nodi, quindi incrementali con num_nodes
+        if key in B_KEYS:
+            return self.num_nodes
+        return super().__inc__(key, value, *args, **kwargs)

ThermalSolver/model/local_gno.py

@@ -1,100 +1,167 @@
 import torch
 from torch import nn
 from torch_geometric.nn import MessagePassing
+from matplotlib.tri import Triangulation
 
 
-# ---- FiLM that starts as identity and normalizes the target ----
+def _import_boundary_conditions(x, boundary, boundary_mask):
-class FiLM(nn.Module):
+    x[boundary_mask] = boundary
-    def __init__(self, c_ch, h_ch):
+
+
+def plot_results_fn(x, pos, i, batch):
+    x = x[batch == 0]
+    pos = pos[batch == 0]
+    tria = Triangulation(pos[:, 0].cpu(), pos[:, 1].cpu())
+    import matplotlib.pyplot as plt
+
+    plt.tricontourf(tria, x[:, 0].cpu(), levels=14)
+    plt.colorbar()
+    plt.savefig(f"out_{i:03d}.png")
+    plt.axis("equal")
+    plt.close()
+
+
+class EncX(nn.Module):
+    def __init__(self, x_ch, hidden):
         super().__init__()
         self.net = nn.Sequential(
-            nn.Linear(c_ch, 2 * h_ch), nn.SiLU(), nn.Linear(2 * h_ch, 2 * h_ch)
+            nn.Linear(x_ch, hidden // 2),
+            nn.SiLU(),
+            nn.Linear(hidden // 2, hidden),
         )
-        # 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)
 
-    def forward(self, h, c):
+    def forward(self, x):
-        gb = self.net(c)
+        return self.net(x)
-        gamma, beta = gb.chunk(2, dim=-1)
+
-        return (1 + gamma) * h + beta
+
+class EncC(nn.Module):
+    def __init__(self, c_ch, hidden):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(c_ch, hidden // 2),
+            nn.SiLU(),
+            nn.Linear(hidden // 2, hidden),
+        )
+
+    def forward(self, c):
+        return self.net(c)
+
+
+class DecX(nn.Module):
+    def __init__(self, hidden, out_ch):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(hidden, hidden // 2),
+            nn.SiLU(),
+            nn.Linear(hidden // 2, out_ch),
+        )
+
+    def forward(self, x):
+        return self.net(x)
 
 
 class ConditionalGNOBlock(MessagePassing):
-    """
+    def __init__(self, hidden_ch, edge_ch=0, aggr="mean"):
-    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="add"):
         super().__init__(aggr=aggr, node_dim=0)
-        # FiLM over the message (per-edge)
+        # self.film_msg = FiLM(c_ch=hidden_ch, h_ch=hidden_ch)
-        self.film_msg = FiLM(c_ch=hidden_ch, h_ch=hidden_ch)
+
         self.edge_attr_net = nn.Sequential(
             nn.Linear(edge_ch, hidden_ch // 2),
             nn.SiLU(),
             nn.Linear(hidden_ch // 2, hidden_ch),
+            nn.Tanh(),
         )
-        self.x_net = nn.Sequential(
+
-            nn.Linear(hidden_ch, hidden_ch * 2),
+        self.msg_proj = nn.Sequential(
+            nn.Linear(hidden_ch, hidden_ch),
             nn.SiLU(),
-            nn.Linear(hidden_ch * 2, hidden_ch),
+            nn.Linear(hidden_ch, hidden_ch),
         )
 
+        self.diff_net = nn.Sequential(
+            nn.Linear(hidden_ch, hidden_ch),
+            nn.SiLU(),
+            nn.Linear(hidden_ch, hidden_ch),
+        )
+
+        self.x_net = nn.Sequential(
+            nn.Linear(hidden_ch, hidden_ch),
+            nn.SiLU(),
+            nn.Linear(hidden_ch, hidden_ch),
+        )
+
+        self.c_ij_net = nn.Sequential(
+            nn.Linear(hidden_ch, hidden_ch),
+            nn.SiLU(),
+            nn.Linear(hidden_ch, hidden_ch),
+            nn.Tanh(),
+        )
+
+        self.balancing = nn.Parameter(torch.tensor(0.0))
+        self.alpha = nn.Parameter(torch.tensor(1.0))
+
     def forward(self, x, c, edge_index, edge_attr=None):
         return self.propagate(edge_index, x=x, c=c, edge_attr=edge_attr)
 
-    def update(self, aggr_out, x):
+    def message(self, x_i, x_j, c_i, c_j, edge_attr):
-        return self.x_net(x) + aggr_out
-
-    def message(self, x_j, c_i, c_j, edge_attr):
-        # c_ij = (c_i + c_j)/2
         c_ij = 0.5 * (c_i + c_j)
-        m = self.film_msg(x_j, c_ij)
+        alpha = torch.sigmoid(self.balancing)
-        if edge_attr is not None:
+        m = alpha * self.diff_net(x_j - x_i) + (1 - alpha) * self.x_net(x_j)
-            a_ij = self.edge_attr_net(edge_attr)
+        m = m * self.c_ij_net(c_ij)
-            m = m * a_ij
+        gate = self.edge_attr_net(edge_attr)
-        return m
+        return m * torch.sigmoid(gate)
+
+    def update(self, aggr_out, x):
+        return x + self.alpha * self.msg_proj(aggr_out)
 
 
 class GatingGNO(nn.Module):
     """
-    In:
+    TODO: add doc
-      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(
+        self.encoder_x = EncX(x_ch_node, hidden)
-            nn.Linear(x_ch_node, hidden // 2),
+        self.encoder_c = EncC(f_ch_node, hidden)
-            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(
+        self.dec = DecX(hidden, out_ch)
-            nn.Linear(hidden, hidden // 2),
-            nn.SiLU(),
-            nn.Linear(hidden // 2, out_ch),
-        )
 
-    def forward(self, x, c, edge_index, edge_attr=None):
+    def forward(
-        x = self.encoder_x(x)  # [N,H]
+        self,
-        c = self.encoder_c(c)  # [N,H]
+        x,
+        c,
+        boundary,
+        boundary_mask,
+        edge_index,
+        edge_attr=None,
+        unrolling_steps=1,
+        plot_results=False,
+        batch=None,
+        pos=None,
+    ):
+        x = self.encoder_x(x)
+        c = self.encoder_c(c)
+        boundary = self.encoder_x(boundary)
+        if plot_results:
+            _import_boundary_conditions(x, boundary, boundary_mask)
+            x_ = self.dec(x)
+            plot_results_fn(x_, pos, 0, batch=batch)
+
+        for _ in range(1, unrolling_steps + 1):
+            _import_boundary_conditions(x, boundary, boundary_mask)
             for blk in self.blocks:
                 x = blk(x, c, edge_index, edge_attr=edge_attr)
+            if plot_results:
+                x_ = self.dec(x)
+                plot_results_fn(x_, pos, _, batch=batch)
+
         return self.dec(x)

ThermalSolver/module.py

@@ -1,6 +1,20 @@
 import torch
 from lightning import LightningModule
 from torch_geometric.data import Batch
+from matplotlib.tri import Triangulation
+
+
+# def plot_results(x, pos, step, i, batch):
+#     x = x[batch == 0]
+#     pos = pos[batch == 0]
+#     tria = Triangulation(pos[:, 0].cpu(), pos[:, 1].cpu())
+#     import matplotlib.pyplot as plt
+
+#     plt.tricontourf(tria, x[:, 0].cpu(), levels=14)
+#     plt.colorbar()
+#     plt.savefig(f"{step:03d}_out_{i:03d}.png")
+#     plt.axis("equal")
+#     plt.close()
 
 
 class GraphSolver(LightningModule):
@@ -8,7 +22,7 @@ class GraphSolver(LightningModule):
         self,
         model: torch.nn.Module,
         loss: torch.nn.Module = None,
-        unrolling_steps: int = 10,
+        unrolling_steps: int = 48,
     ):
         super().__init__()
         self.model = model
@@ -19,13 +33,21 @@ class GraphSolver(LightningModule):
         self,
         x: torch.Tensor,
         c: torch.Tensor,
+        boundary: torch.Tensor,
+        boundary_mask: torch.Tensor,
         edge_index: torch.Tensor,
         edge_attr: torch.Tensor,
+        unrolling_steps: int = None,
     ):
-        return self.model(x, c, edge_index, edge_attr)
+        return self.model(
-
+            x,
-    def _compute_loss_train(self, x, x_prev, y):
+            c,
-        return self.loss(x, y) + self.loss(x, x_prev)
+            boundary,
+            boundary_mask,
+            edge_index,
+            edge_attr,
+            unrolling_steps,
+        )
 
     def _compute_loss(self, x, y):
         return self.loss(x, y)
@@ -46,35 +68,55 @@ class GraphSolver(LightningModule):
 
     def training_step(self, batch: Batch, _):
         x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
-        loss = 0.0
+        # x = self._impose_bc(x, batch)
-        for _ in range(self.unrolling_steps):
+        # for _ in range(self.unrolling_steps):
-            x_prev = x.detach()
+        y_pred = self(
-            x = self(x_prev, c, edge_index=edge_index, edge_attr=edge_attr)
+            x,
-            actual_loss = self.loss(x, y)
+            c,
-            loss += actual_loss
+            batch.boundary_values,
-            print(f"Train step loss: {actual_loss.item()}")
+            batch.boundary_mask,
-
+            edge_index=edge_index,
+            edge_attr=edge_attr,
+            unrolling_steps=self.unrolling_steps,
+        )
+        # x = self._impose_bc(x, batch)
+        loss = self.loss(y_pred, y)
         self._log_loss(loss, batch, "train")
         return loss
 
     def validation_step(self, batch: Batch, _):
         x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
-        for _ in range(self.unrolling_steps):
+        y_pred = self(
-            x_prev = x.detach()
+            x,
-            x = self(x_prev, c, edge_index=edge_index, edge_attr=edge_attr)
+            c,
-            loss = self.loss(x, x_prev)
+            batch.boundary_values,
-            if loss < 1e-5:
+            batch.boundary_mask,
-                break
+            edge_index=edge_index,
-        loss = self._compute_loss(x, y)
+            edge_attr=edge_attr,
+            unrolling_steps=self.unrolling_steps,
+        )
+        loss = self.loss(y_pred, y)
         self._log_loss(loss, batch, "val")
         return loss
 
     def test_step(self, batch: Batch, _):
         x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
-        for _ in range(self.unrolling_steps):
+        # for _ in range(self.unrolling_steps):
-            x_prev = x.detach()
+        y_pred = self.model(
-            x = self(x_prev, c, edge_index=edge_index, edge_attr=edge_attr)
+            x,
-        loss = self._compute_loss(x, y)
+            c,
+            batch.boundary_values,
+            batch.boundary_mask,
+            edge_index=edge_index,
+            edge_attr=edge_attr,
+            unrolling_steps=self.unrolling_steps,
+            plot_results=True,
+            batch=batch.batch,
+            pos=batch.pos,
+        )
+        # x = self._impose_bc(x, batch)
+        # plot_results(x, batch.pos, self.global_step, _, batch.batch)
+        loss = self._compute_loss(y_pred, y)
         self._log_loss(loss, batch, "test")
         return loss
 
@@ -82,6 +124,6 @@ class GraphSolver(LightningModule):
         optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
         return optimizer
 
-    def scale_bc(self, data: Batch, y: torch.Tensor):
+    def _impose_bc(self, x: torch.Tensor, data: Batch):
-        t = data.boundary_temperature[data.batch]
+        x[data.boundary_mask] = data.boundary_values
-        return y * t
+        return x

ThermalSolver/normalizer.py

@@ -0,0 +1,51 @@
+import torch
+from torch_geometric.data import Data
+
+D_IN_KEYS = "x"
+D_ATTR_KEYS = ["c", "edge_attr"]
+D_OUT_KEY = "y"
+D_KEYS = [D_IN_KEYS] + [D_OUT_KEY] + D_ATTR_KEYS
+D_BOUNDS_KEYS = "boundary_temperatures"
+
+
+class Normalizer:
+    def __init__(self, data):
+        self.mean, self.std = self._compute_stats(data)
+
+    def _compute_stats(self, data: list[Data]) -> tuple[dict, dict]:
+        mean = {}
+        std = {}
+        for key in D_KEYS:
+            tmp = torch.empty(0)
+            for d in data:
+                if not hasattr(d, key):
+                    raise AttributeError(f"Manca '{key}' in uno dei Data.")
+                if tmp.numel() == 0:
+                    tmp = d[key]
+                else:
+                    tmp = torch.cat([tmp, d[key]], dim=0)
+            mean[key] = tmp.mean(dim=0, keepdim=True)
+            std[key] = tmp.std(dim=0, keepdim=True) + 1e-6
+        return mean, std
+
+    def normalize(self, data):
+        for d in data:
+            for key in D_KEYS:
+                if not hasattr(d, key):
+                    raise AttributeError(f"Manca '{key}' in uno dei Data.")
+                d[key] = (d[key] - self.mean[key]) / self.std[key]
+        self._recompute_boundary_temperatures(data)
+
+    def _recompute_boundary_temperatures(self, data):
+        for d in data:
+            bottom_bc = d.y[d.bottom_boundary_ids].median()
+            top_bc = d.y[d.top_boundary_ids].median()
+            left_bc = d.y[d.left_boundary_ids].median()
+            right_bc = d.y[d.right_boundary_ids].median()
+            boundaries_temperatures = torch.tensor(
+                [bottom_bc, right_bc, top_bc, left_bc], dtype=torch.float32
+            )
+            d.boundary_temperatures = boundaries_temperatures.unsqueeze(0)
+
+    def denormalize(self, y: torch.tensor):
+        return y * self.std[D_OUT_KEY] + self.mean[D_OUT_KEY]