260 lines
7.8 KiB
Python
260 lines
7.8 KiB
Python
import torch
|
|
from torch import nn
|
|
from torch_geometric.nn import MessagePassing
|
|
from matplotlib.tri import Triangulation
|
|
|
|
|
|
def plot_results_fn(x, pos, i, batch):
|
|
x = x[batch == 0]
|
|
pos = pos[batch == 0]
|
|
tria = Triangulation(pos[:, 0].cpu(), pos[:, 1].cpu())
|
|
import matplotlib.pyplot as plt
|
|
|
|
plt.tricontourf(tria, x[:, 0].cpu(), levels=14)
|
|
plt.colorbar()
|
|
plt.savefig(f"out_{i:03d}.png")
|
|
plt.axis("equal")
|
|
plt.close()
|
|
|
|
|
|
class EncX(nn.Module):
|
|
def __init__(self, x_ch, hidden):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Linear(x_ch, hidden // 2),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden // 2, hidden),
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|
|
|
|
|
|
class EncC(nn.Module):
|
|
def __init__(self, c_ch, hidden):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Linear(c_ch, hidden // 2),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden // 2, hidden),
|
|
)
|
|
|
|
def forward(self, c):
|
|
return self.net(c)
|
|
|
|
|
|
class DecX(nn.Module):
|
|
def __init__(self, hidden, out_ch):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Linear(hidden, hidden // 2),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden // 2, out_ch),
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|
|
|
|
|
|
# class ConditionalGNOBlock(MessagePassing):
|
|
# def __init__(self, hidden_ch, edge_ch=0, aggr="mean"):
|
|
# super().__init__(aggr=aggr, node_dim=0)
|
|
|
|
# self.edge_attr_net = nn.Sequential(
|
|
# nn.Linear(edge_ch, hidden_ch // 2),
|
|
# nn.SiLU(),
|
|
# nn.Linear(hidden_ch // 2, 1),
|
|
# nn.Softplus()
|
|
# )
|
|
|
|
# self.diff_net = nn.Sequential(
|
|
# nn.Linear(hidden_ch, hidden_ch),
|
|
# nn.SiLU(),
|
|
# nn.Linear(hidden_ch, hidden_ch),
|
|
# )
|
|
|
|
# # self.x_net = nn.Sequential(
|
|
# # nn.Linear(hidden_ch, hidden_ch),
|
|
# # nn.SiLU(),
|
|
# # nn.Linear(hidden_ch, hidden_ch),
|
|
# # )
|
|
|
|
# self.c_ij_net = nn.Sequential(
|
|
# nn.Linear(hidden_ch, hidden_ch // 2),
|
|
# nn.SiLU(),
|
|
# nn.Linear(hidden_ch // 2, 1),
|
|
# nn.Sigmoid(),
|
|
# )
|
|
|
|
# # self.gamma_net = nn.Sequential(
|
|
# # nn.Linear(2 * hidden_ch, hidden_ch),
|
|
# # nn.SiLU(),
|
|
# # nn.Linear(hidden_ch, hidden_ch // 2),
|
|
# # nn.SiLU(),
|
|
# # nn.Linear(hidden_ch // 2, 1),
|
|
# # nn.Sigmoid(),
|
|
# # )
|
|
|
|
# self.alpha_net = nn.Sequential(
|
|
# nn.Linear(2 * hidden_ch, hidden_ch),
|
|
# nn.SiLU(),
|
|
# nn.Linear(hidden_ch, hidden_ch // 2),
|
|
# nn.SiLU(),
|
|
# nn.Linear(hidden_ch // 2, 1),
|
|
# nn.Sigmoid(),
|
|
# )
|
|
|
|
# def forward(self, x, c, edge_index, edge_attr=None):
|
|
# return self.propagate(edge_index, x=x, c=c, edge_attr=edge_attr)
|
|
|
|
# def message(self, x_i, x_j, c_i, c_j, edge_attr):
|
|
# c_ij = 0.5 * (c_i + c_j)
|
|
# # gamma = self.gamma_net(torch.cat([x_i, x_j], dim=-1))
|
|
# # gate = torch.sself.edge_attr_net(edge_attr))
|
|
# gate = self.edge_attr_net(edge_attr)
|
|
# # m = (
|
|
# # gamma * self.diff_net(x_j - x_i) + (1 - gamma) * self.x_net(x_j)
|
|
# # ) * gate
|
|
# m = self.diff_net(x_j - x_i) * gate
|
|
# m = m * self.c_ij_net(c_ij)
|
|
# return m
|
|
|
|
# def update(self, aggr_out, x):
|
|
# alpha = self.alpha_net(torch.cat([x, aggr_out], dim=-1))
|
|
# return x + alpha * aggr_out
|
|
|
|
|
|
class ConditionalGNOBlock(MessagePassing):
|
|
def __init__(self, hidden_ch, edge_ch=0, aggr="mean"):
|
|
super().__init__(aggr=aggr, node_dim=0)
|
|
self.edge_ch = edge_ch
|
|
|
|
# Rete che mappa edge_attr -> coefficiente scalare (log-scale)
|
|
# Se edge_ch==0 useremo un coefficiente apprendibile globale
|
|
self.edge_attr_net = nn.Sequential(
|
|
nn.Linear(edge_ch, hidden_ch),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_ch, hidden_ch // 2),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_ch // 2, 1),
|
|
nn.Softplus(),
|
|
)
|
|
# gating dalla condizione c_ij (restituisce scalar in (0,1))
|
|
self.c_ij_net = nn.Sequential(
|
|
nn.Linear(hidden_ch, hidden_ch),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_ch, hidden_ch // 2),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_ch // 2, 1),
|
|
nn.Sigmoid(),
|
|
)
|
|
|
|
# alpha per passo (clampato tramite sigmoid)
|
|
self.alpha_net = nn.Sequential(
|
|
nn.Linear(2 * hidden_ch, hidden_ch),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_ch, hidden_ch // 2),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_ch // 2, 1),
|
|
nn.Sigmoid(),
|
|
)
|
|
|
|
# self.norm = nn.LayerNorm(hidden_ch)
|
|
|
|
def forward(self, x, c, edge_index, edge_attr=None):
|
|
# chiamiamo propagate; edge_attr può essere None
|
|
return self.propagate(edge_index, x=x, c=c, edge_attr=edge_attr)
|
|
|
|
def message(self, x_i, x_j, c_i, c_j, edge_attr):
|
|
"""
|
|
Implementazione diffusiva:
|
|
m_ij = w_ij * (x_j - x_i) * c_gate_ij
|
|
dove w_ij = softplus(edge_attr_net(edge_attr)) >= 0
|
|
"""
|
|
# 1) calcola c_ij e gating da c
|
|
c_ij = 0.5 * (c_i + c_j) # [E, H]
|
|
c_gate = self.c_ij_net(c_ij) # [E, 1] in (0,1)
|
|
|
|
# 2) calcola peso scalare non-negativo per edge
|
|
w_raw = self.edge_attr_net(edge_attr) # [E,1]
|
|
|
|
# softplus -> peso >= 0; aggiungo epsilon per stabilità
|
|
w = w_raw + 1e-12 # [E,1]
|
|
|
|
# 3) messaggio base: differenza pesata
|
|
diff = x_j - x_i # [E, H]
|
|
m = w * diff # broadcast: [E,1] * [E,H] -> [E,H]
|
|
|
|
# 4) applica gating dalla condizione
|
|
m = m * c_gate # [E,H]
|
|
|
|
# Restituisco anche w (sfruttabile in update) — ma MessagePassing non ritorna extra,
|
|
# così se vuoi degree-normalization devi calcolare i gradi prima di propagate.
|
|
# Qui ritorno solo m: la normalizzazione per grado la faccio in update usando 'mean' aggr
|
|
return m
|
|
|
|
def update(self, aggr_out, x):
|
|
"""
|
|
aggr_out:
|
|
- se aggr='sum': somma delle w_ij*(x_j-x_i) incoming
|
|
- se aggr='mean': già normalizzato sul numero di vicini (ma non per somma dei pesi)
|
|
Qui normalizziamo implicitamente dividendo per (1 + |aggr_out|_norm) per stabilità,
|
|
e applichiamo il passo alpha.
|
|
"""
|
|
# aggr_out = self.norm(aggr_out) # stabilizza la scala
|
|
|
|
# alpha vettoriale/scalar: [N,1]
|
|
alpha = self.alpha_net(torch.cat([x, aggr_out], dim=-1)) # in (0,1)
|
|
|
|
x_new = x + alpha * aggr_out
|
|
|
|
return x_new
|
|
|
|
|
|
class GatingGNO(nn.Module):
|
|
"""
|
|
TODO: add doc
|
|
"""
|
|
|
|
def __init__(
|
|
self, x_ch_node, f_ch_node, hidden, layers, edge_ch=0, out_ch=1
|
|
):
|
|
super().__init__()
|
|
self.encoder_x = EncX(x_ch_node, hidden)
|
|
self.encoder_c = EncC(f_ch_node, hidden)
|
|
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
ConditionalGNOBlock(hidden_ch=hidden, edge_ch=edge_ch)
|
|
for _ in range(layers)
|
|
]
|
|
)
|
|
self.dec = DecX(hidden, out_ch)
|
|
|
|
def forward(
|
|
self,
|
|
x,
|
|
c,
|
|
boundary,
|
|
boundary_mask,
|
|
edge_index,
|
|
edge_attr=None,
|
|
unrolling_steps=1,
|
|
plot_results=False,
|
|
batch=None,
|
|
pos=None,
|
|
):
|
|
x = self.encoder_x(x)
|
|
c = self.encoder_c(c)
|
|
if plot_results:
|
|
x_ = self.dec(x)
|
|
plot_results_fn(x_, pos, 0, batch=batch)
|
|
for _ in range(1, unrolling_steps + 1):
|
|
for blk in self.blocks:
|
|
x = blk(x, c, edge_index, edge_attr=edge_attr)
|
|
if plot_results:
|
|
x_ = self.dec(x)
|
|
plot_results_fn(x_, pos, _, batch=batch)
|
|
|
|
return self.dec(x)
|