diff --git a/ThermalSolver/data_module.py b/ThermalSolver/data_module.py index f57590f..127fcea 100644 --- a/ThermalSolver/data_module.py +++ b/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), - torch.tensor(snapshot["boundary_values"], dtype=torch.float32), - torch.tensor(snapshot["temperature"], dtype=torch.float32), + snapshot, 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, - boundary_vales: torch.Tensor, - temperature: torch.Tensor, + snapshot: dict, 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 - return Data( - x=boundary_vales.unsqueeze(-1), + + left_ids = left_ids[~torch.isin(left_ids, bottom_ids)] + right_ids = right_ids[~torch.isin(right_ids, bottom_ids)] + 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: - return DataLoader(self.val_data, batch_size=self.batch_size) + def val_dataloader(self): + return DataLoader( + self.val_data, batch_size=self.batch_size, shuffle=False + ) - def test_dataloader(self) -> DataLoader: - return DataLoader(self.test_data, batch_size=self.batch_size) + def test_dataloader(self): + return DataLoader( + self.test_data, batch_size=self.batch_size, shuffle=False + ) diff --git a/ThermalSolver/mesh_data.py b/ThermalSolver/mesh_data.py new file mode 100644 index 0000000..1c32b05 --- /dev/null +++ b/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) diff --git a/ThermalSolver/model/local_gno.py b/ThermalSolver/model/local_gno.py index 7bd58f1..6432850 100644 --- a/ThermalSolver/model/local_gno.py +++ b/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 ---- -class FiLM(nn.Module): - def __init__(self, c_ch, h_ch): +def _import_boundary_conditions(x, boundary, boundary_mask): + x[boundary_mask] = boundary + + +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): - gb = self.net(c) - gamma, beta = gb.chunk(2, dim=-1) - return (1 + gamma) * h + beta + def forward(self, x): + return self.net(x) + + +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): - """ - 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"): + def __init__(self, hidden_ch, edge_ch=0, aggr="mean"): 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): - 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 + def message(self, x_i, x_j, c_i, c_j, edge_attr): c_ij = 0.5 * (c_i + c_j) - m = self.film_msg(x_j, c_ij) - if edge_attr is not None: - a_ij = self.edge_attr_net(edge_attr) - m = m * a_ij - return m + alpha = torch.sigmoid(self.balancing) + m = alpha * self.diff_net(x_j - x_i) + (1 - alpha) * self.x_net(x_j) + m = m * self.c_ij_net(c_ij) + gate = self.edge_attr_net(edge_attr) + 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: - x : [N, Cx] (e.g., u or features to predict from) - c : [N, Cf] (conditioning field, e.g., conductivity) - Out: - y : [N, out_ch] + TODO: add doc """ 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.encoder_x = EncX(x_ch_node, hidden) + self.encoder_c = EncC(f_ch_node, hidden) + self.blocks = nn.ModuleList( [ ConditionalGNOBlock(hidden_ch=hidden, edge_ch=edge_ch) for _ in range(layers) ] ) - self.dec = nn.Sequential( - nn.Linear(hidden, hidden // 2), - nn.SiLU(), - nn.Linear(hidden // 2, out_ch), - ) + self.dec = DecX(hidden, out_ch) + + def forward( + self, + 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) - 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) diff --git a/ThermalSolver/module.py b/ThermalSolver/module.py index 2acd94f..f9d9e2e 100644 --- a/ThermalSolver/module.py +++ b/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) - - def _compute_loss_train(self, x, x_prev, y): - return self.loss(x, y) + self.loss(x, x_prev) + return self.model( + x, + c, + 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 - for _ in range(self.unrolling_steps): - x_prev = x.detach() - x = self(x_prev, c, edge_index=edge_index, edge_attr=edge_attr) - actual_loss = self.loss(x, y) - loss += actual_loss - print(f"Train step loss: {actual_loss.item()}") - + # x = self._impose_bc(x, batch) + # for _ in range(self.unrolling_steps): + y_pred = self( + x, + c, + batch.boundary_values, + 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): - x_prev = x.detach() - x = self(x_prev, c, edge_index=edge_index, edge_attr=edge_attr) - loss = self.loss(x, x_prev) - if loss < 1e-5: - break - loss = self._compute_loss(x, y) + y_pred = self( + x, + c, + batch.boundary_values, + batch.boundary_mask, + edge_index=edge_index, + 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): - x_prev = x.detach() - x = self(x_prev, c, edge_index=edge_index, edge_attr=edge_attr) - loss = self._compute_loss(x, y) + # for _ in range(self.unrolling_steps): + y_pred = self.model( + x, + 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): - t = data.boundary_temperature[data.batch] - return y * t + def _impose_bc(self, x: torch.Tensor, data: Batch): + x[data.boundary_mask] = data.boundary_values + return x diff --git a/ThermalSolver/normalizer.py b/ThermalSolver/normalizer.py new file mode 100644 index 0000000..4a8b62e --- /dev/null +++ b/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]