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

fix model

Browse Source
This commit is contained in:
Filippo Olivo
2025-11-14 17:06:08 +01:00
parent e1117d89c6
commit 94ad6ff160
1 changed files with 6 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
ThermalSolver/model/learnable_finite_difference.py
View File

@@ -9,7 +9,8 @@ class FiniteDifferenceStep(MessagePassing):
TODO: add docstring. TODO: add docstring.
""" """
def __init__(self, edge_ch=5, hidden_dim=16, aggr: str = "add"): def __init__(self, hidden_dim=16, aggr: str = "add"):
print(aggr)
super().__init__(aggr=aggr) super().__init__(aggr=aggr)
self.x_embedding = nn.Sequential( self.x_embedding = nn.Sequential(
spectral_norm(nn.Linear(1, hidden_dim // 2)), spectral_norm(nn.Linear(1, hidden_dim // 2)),
@@ -17,12 +18,6 @@ class FiniteDifferenceStep(MessagePassing):
spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)), spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)),
) )
# self.update_net = nn.Sequential(
# spectral_norm(nn.Linear(2 * hidden_dim, hidden_dim)),
# nn.GELU(),
# spectral_norm(nn.Linear(hidden_dim, hidden_dim)),
# )
self.out_net = nn.Sequential( self.out_net = nn.Sequential(
spectral_norm(nn.Linear(hidden_dim, hidden_dim // 2)), spectral_norm(nn.Linear(hidden_dim, hidden_dim // 2)),
nn.GELU(), nn.GELU(),
@@ -35,20 +30,18 @@ class FiniteDifferenceStep(MessagePassing):
""" """
x_ = self.x_embedding(x) x_ = self.x_embedding(x)
out = self.propagate(edge_index, x=x_, edge_attr=edge_attr, deg=deg) out = self.propagate(edge_index, x=x_, edge_attr=edge_attr, deg=deg)
return self.out_net(x_ + out) return self.out_net(out)
def message(self, x_i, x_j, edge_attr): def message(self, x_j, edge_attr):
""" """
TODO: add docstring. TODO: add docstring.
""" """
return (x_j - x_i) * edge_attr.view(-1, 1) return x_j * edge_attr.view(-1, 1)
def update(self, aggr_out, x): def update(self, aggr_out, _):
""" """
TODO: add docstring. TODO: add docstring.
""" """
# update_input = torch.cat([x, aggr_out], dim=-1)
# return self.update_net(update_input)
return aggr_out return aggr_out
def aggregate(self, inputs, index, deg): def aggregate(self, inputs, index, deg):