fix model and datamodule

ThermalSolver/graph_datamodule_unsteady.py

@@ -1,18 +1,19 @@
 import torch
 from tqdm import tqdm
 from lightning import LightningDataModule
-from datasets import load_dataset
+from datasets import load_dataset, concatenate_datasets
 from torch_geometric.data import Data
 from torch_geometric.loader import DataLoader
 from torch_geometric.utils import to_undirected
 from .mesh_data import MeshData
+from typing import List, Union
 
 
 class GraphDataModule(LightningDataModule):
     def __init__(
         self,
         hf_repo: str,
-        split_name: str,
+        split_name: Union[str, List[str]],
         n_elements: int = None,
         train_size: float = 0.2,
         val_size: float = 0.1,
@@ -44,8 +45,30 @@ class GraphDataModule(LightningDataModule):
         self.radius = radius
 
     def prepare_data(self):
-        dataset = load_dataset(self.hf_repo, name="snapshots")[self.split_name]
-        geometry = load_dataset(self.hf_repo, name="geometry")[self.split_name]
+        if isinstance(self.split_name, list):
+            dataset_list = []
+            geometry_list = []
+            for split in self.split_name:
+                dataset_list.append(
+                    load_dataset(self.hf_repo, name="snapshots")[split]
+                )
+                geometry_list.append(
+                    load_dataset(self.hf_repo, name="geometry")[split]
+                )
+
+            dataset = concatenate_datasets(dataset_list)
+            geometry = concatenate_datasets(geometry_list)
+            idx = torch.randperm(len(dataset))
+            dataset = dataset.select(idx.tolist())
+            geometry = geometry.select(idx.tolist())
+        else:
+            dataset = load_dataset(self.hf_repo, name="snapshots")[
+                self.split_name
+            ]
+            geometry = load_dataset(self.hf_repo, name="geometry")[
+                self.split_name
+            ]
+
         if self.n_elements is not None:
             dataset = dataset.select(range(self.n_elements))
             geometry = geometry.select(range(self.n_elements))
@@ -86,7 +109,7 @@ class GraphDataModule(LightningDataModule):
                 dim=0,
             )
         )
-        print(temperatures.shape)
+        # print(temperatures.shape)
 
         pos = torch.tensor(geometry["points"], dtype=torch.float32)[:, :2]
 
@@ -103,9 +126,7 @@ class GraphDataModule(LightningDataModule):
         boundary_mask = torch.tensor(
             geometry["constraints_mask"], dtype=torch.int64
         )
-        boundary_values = torch.tensor(
-            geometry["constraints_values"], dtype=torch.float32
-        )
+        boundary_values = temperatures[0, boundary_mask]
 
         edge_attr = torch.norm(pos[edge_index[0]] - pos[edge_index[1]], dim=1)
         if self.remove_boundary_edges: