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

try a new model

Browse Source
This commit is contained in:
Filippo Olivo
2025-11-12 15:20:43 +01:00
parent a2dd348423
commit dc59114f4a
3 changed files with 31 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
ThermalSolver/graph_module.py
View File

@@ -75,9 +75,6 @@ class GraphSolver(LightningModule):
def _compute_loss(self, x, y):
return self.loss(x, y)
def _preprocess_batch(self, batch: Batch):
return batch.x, batch.y, batch.c, batch.edge_index, batch.edge_attr
def _log_loss(self, loss, batch, stage: str):
self.log(
f"{stage}/loss",
@@ -115,19 +112,25 @@ class GraphSolver(LightningModule):
self.manual_backward(loss / max_acc_iters)
return loss_.item()
def _preprocess_batch(self, batch: Batch):
x, y, c, edge_index, edge_attr = (
batch.x,
batch.y,
batch.c,
batch.edge_index,
batch.edge_attr,
)
edge_attr = 1 / edge_attr
c_ij = self._compute_c_ij(c, edge_index)
edge_attr = edge_attr * c_ij
return x, y, edge_index, edge_attr
def training_step(self, batch: Batch, batch_idx: int):
optim = self.optimizers()
optim.zero_grad()
x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
x, y, edge_index, edge_attr = self._preprocess_batch(batch)
edge_w = 1 / edge_attr[:, -1]
c_ij = self._compute_c_ij(c, edge_index)
edge_w = edge_w * c_ij
deg = self._compute_deg(edge_index, edge_w, x.size(0))
edge_attr = torch.cat(
[edge_attr, edge_w.unsqueeze(-1), c_ij.unsqueeze(-1)], dim=1
)
deg = self._compute_deg(edge_index, edge_attr, x.size(0))
losses = []
acc_loss, acc_it = 0, 0
max_acc_iters = (
@@ -139,7 +142,7 @@ class GraphSolver(LightningModule):
out = self._compute_model_steps(
x,
edge_index,
edge_attr,
edge_attr.unsqueeze(-1),
deg,
batch.boundary_mask,
batch.boundary_values,
@@ -199,21 +202,15 @@ class GraphSolver(LightningModule):
return super().on_train_epoch_end()
def validation_step(self, batch: Batch, _):
x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
x, y, edge_index, edge_attr = self._preprocess_batch(batch)
edge_w = 1 / edge_attr[:, -1]
c_ij = self._compute_c_ij(c, edge_index)
edge_w = edge_w * c_ij
deg = self._compute_deg(edge_index, edge_w, x.size(0))
deg = self._compute_deg(edge_index, edge_attr, x.size(0))
edge_attr = torch.cat(
[edge_attr, edge_w.unsqueeze(-1), c_ij.unsqueeze(-1)], dim=1
)
for i in range(self.current_iters):
out = self._compute_model_steps(
x,
edge_index,
edge_attr,
edge_attr.unsqueeze(-1),
deg,
batch.boundary_mask,
batch.boundary_values,
@@ -252,20 +249,15 @@ class GraphSolver(LightningModule):
# loss = self._compute_loss(y_pred, y)
# # _plot_mesh(batch.pos, y, y_pred, batch.batch)
# self._log_loss(loss, batch, "test")
x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
edge_w = 1 / edge_attr[:, -1]
c_ij = self._compute_c_ij(c, edge_index)
edge_w = edge_w * c_ij
deg = self._compute_deg(edge_index, edge_w, x.size(0))
x, y, edge_index, edge_attr = self._preprocess_batch(batch)
deg = self._compute_deg(edge_index, edge_attr, x.size(0))
edge_attr = torch.cat(
[edge_attr, edge_w.unsqueeze(-1), c_ij.unsqueeze(-1)], dim=1
)
for i in range(self.max_iters):
out = self._compute_model_steps(
x,
edge_index,
edge_attr,
edge_attr.unsqueeze(-1),
deg,
batch.boundary_mask,
batch.boundary_values,
@@ -278,7 +270,6 @@ class GraphSolver(LightningModule):
loss = self.loss(out, y)
self._log_loss(loss, batch, "test")
x = u
self.log(
"test/iterations",
i + 1,