folivo/thermal-conduction-ml
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

implement new model

Browse Source
This commit is contained in:
Filippo Olivo
2025-10-03 10:42:05 +02:00
parent b1a9cecb42
commit 7a6fbdb89c
3 changed files with 80 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
ThermalSolver/data_module.py
View File

@@ -17,6 +17,7 @@ class GraphDataModule(LightningDataModule):
val_size: float = 0.1,
test_size: float = 0.1,
batch_size: int = 32,
remove_boundary_edges: bool = True,
):
super().__init__()
self.hf_repo = hf_repo
@@ -27,6 +28,7 @@ class GraphDataModule(LightningDataModule):
self.val_size = val_size
self.test_size = test_size
self.batch_size = batch_size
self.remove_boundary_edges = remove_boundary_edges
def prepare_data(self):
hf_dataset = load_dataset(self.hf_repo, name="snapshots")[
@@ -44,6 +46,28 @@ class GraphDataModule(LightningDataModule):
)
]
def _compute_boundary_mask(
self, bottom_ids, right_ids, top_ids, left_ids, temperature
):
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 boundary_mask, boundary_values
def _build_dataset(
self,
snapshot: dict,
@@ -66,27 +90,34 @@ class GraphDataModule(LightningDataModule):
)
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 = torch.cat(
[edge_attr, torch.norm(edge_attr, dim=1).unsqueeze(-1)], dim=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)
x = torch.zeros_like(temperature, dtype=torch.float32).unsqueeze(-1)
if self.remove_boundary_edges:
x[boundary_mask] = boundary_values.unsqueeze(-1)
return MeshData(
x=x,
c=conductivity.unsqueeze(-1),
edge_index=edge_index,
pos=pos,
edge_attr=edge_attr,
y=temperature.unsqueeze(-1),
boundary_mask=boundary_mask,
boundary_values=torch.tensor(0),
)
return MeshData(
x=torch.rand_like(temperature).unsqueeze(-1),
@@ -110,7 +141,6 @@ class GraphDataModule(LightningDataModule):
if stage == "test" or stage is None:
self.test_data = self.data[val_end:]
# nel tuo LightningDataModule
def train_dataloader(self):
return DataLoader(
self.train_data, batch_size=self.batch_size, shuffle=True