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

new data format

Browse Source
This commit is contained in:
Filippo Olivo
2025-12-12 10:18:16 +01:00
parent 27c2aeb736
commit 732d48c360
3 changed files with 118 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
ThermalSolver/graph_datamodule_unsteady.py
View File

@@ -64,28 +64,6 @@ class GraphDataModule(LightningDataModule):
"test": geometry.select(range(train_len + valid_len, total_len)),
}
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,
@@ -96,25 +74,22 @@ class GraphDataModule(LightningDataModule):
geometry["conductivity"], dtype=torch.float32
)
temperatures = (
torch.tensor(snapshot["temperatures"], dtype=torch.float32)[:40]
torch.tensor(snapshot["unsteady"], dtype=torch.float32)
if not test
else torch.tensor(snapshot["temperatures"], dtype=torch.float32)[
: self.unrolling_steps + 1
]
else torch.stack(
[
torch.tensor(snapshot["unsteady"], dtype=torch.float32)[
0, ...
],
torch.tensor(snapshot["steady"], dtype=torch.float32),
],
dim=0,
)
)
times = torch.tensor(snapshot["times"], dtype=torch.float32)
print(temperatures.shape)
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
)
if self.build_radial_graph:
raise NotImplementedError(
"Radial graph building not implemented yet."
@@ -125,17 +100,37 @@ class GraphDataModule(LightningDataModule):
).T
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, temperatures[0, :]
boundary_mask = torch.tensor(
geometry["constraints_mask"], dtype=torch.int64
)
boundary_values = torch.tensor(
geometry["constraints_values"], dtype=torch.float32
)
edge_attr = torch.norm(pos[edge_index[0]] - pos[edge_index[1]], dim=1)
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 = edge_attr[edge_index_mask]
n_data = temperatures.size(0) - self.unrolling_steps
n_data = max(temperatures.size(0) - self.unrolling_steps, 1)
data = []
if test:
data.append(
MeshData(
x=temperatures[0, :].unsqueeze(-1),
y=temperatures[1:2, :].unsqueeze(-1).permute(1, 0, 2),
c=conductivity.unsqueeze(-1),
edge_index=edge_index,
pos=pos,
edge_attr=edge_attr,
boundary_mask=boundary_mask,
boundary_values=boundary_values,
)
)
return data
for i in range(n_data):
x = temperatures[i, :].unsqueeze(-1)
y = (