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),
-                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
+        )