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

add refined geometries in datamodule

Browse Source
This commit is contained in:
Filippo Olivo
2025-10-02 12:20:22 +02:00
parent 1be98591d4
commit b26403189f
2 changed files with 23 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
ThermalSolver/data_module.py
View File

@@ -35,53 +35,36 @@ class GraphDataModule(LightningDataModule):
self.geometry = load_dataset(self.hf_repo, name="geometry")[
self.split_name
]
edge_index = torch.tensor(
self.geometry["edge_index"][0], dtype=torch.int64
)
pos = torch.tensor(self.geometry["points"][0], dtype=torch.float32)[
:, :2
]
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(
snapshot,
edge_index.T,
pos,
bottom_ids,
top_ids,
left_ids,
right_ids,
self._build_dataset(snapshot, geometry)
for snapshot, geometry in tqdm(
zip(hf_dataset, self.geometry),
desc="Building graphs",
total=len(hf_dataset),
)
for snapshot in tqdm(hf_dataset, desc="Building graphs")
]
def _build_dataset(
self,
snapshot: dict,
edge_index: torch.Tensor,
pos: torch.Tensor,
bottom_ids: torch.Tensor,
top_ids: torch.Tensor,
left_ids: torch.Tensor,
right_ids: torch.Tensor,
geometry: dict,
) -> Data:
conductivity = torch.tensor(
snapshot["conductivity"], dtype=torch.float32
)
temperature = torch.tensor(snapshot["temperature"], dtype=torch.float32)
edge_index = torch.tensor(geometry["edge_index"], dtype=torch.int64).T
pos = torch.tensor(geometry["points"], dtype=torch.float32)[:, :2]
bottom_ids = torch.tensor(
geometry["bottom_boundary_ids"], dtype=torch.long
)
top_ids = torch.tensor(geometry["top_boundary_ids"], dtype=torch.long)
left_ids = torch.tensor(geometry["left_boundary_ids"], dtype=torch.long)
right_ids = torch.tensor(
geometry["right_boundary_ids"], dtype=torch.long
)
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(

9
ThermalSolver/model/local_gno.py
View File

@@ -106,10 +106,13 @@ class ConditionalGNOBlock(MessagePassing):
def message(self, x_i, x_j, c_i, c_j, edge_attr):
c_ij = 0.5 * (c_i + c_j)
alpha = torch.sigmoid(self.balancing)
m = alpha * self.diff_net(x_j - x_i) + (1 - alpha) * self.x_net(x_j)
gate = torch.sigmoid(self.edge_attr_net(edge_attr))
m = (
alpha * self.diff_net(x_j - x_i)
+ (1 - alpha) * self.x_net(x_j) * gate
)
m = m * self.c_ij_net(c_ij)
gate = self.edge_attr_net(edge_attr)
return m * torch.sigmoid(gate)
return m
def update(self, aggr_out, x):
return x + self.alpha * self.msg_proj(aggr_out)