54 lines
1.4 KiB
Python
54 lines
1.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch_geometric.nn import MessagePassing
|
|
from torch.nn.utils import spectral_norm
|
|
|
|
|
|
class FiniteDifferenceStep(MessagePassing):
|
|
"""
|
|
TODO: add docstring.
|
|
"""
|
|
|
|
def __init__(self, hidden_dim=16, aggr: str = "add"):
|
|
print(aggr)
|
|
super().__init__(aggr=aggr)
|
|
self.x_embedding = nn.Sequential(
|
|
spectral_norm(nn.Linear(1, hidden_dim // 2)),
|
|
nn.GELU(),
|
|
spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)),
|
|
)
|
|
|
|
self.out_net = nn.Sequential(
|
|
spectral_norm(nn.Linear(hidden_dim, hidden_dim // 2)),
|
|
nn.GELU(),
|
|
spectral_norm(nn.Linear(hidden_dim // 2, 1)),
|
|
)
|
|
|
|
def forward(self, x, edge_index, edge_attr, deg):
|
|
"""
|
|
TODO: add docstring.
|
|
"""
|
|
x_ = self.x_embedding(x)
|
|
out = self.propagate(edge_index, x=x_, edge_attr=edge_attr, deg=deg)
|
|
return self.out_net(out)
|
|
|
|
def message(self, x_j, edge_attr):
|
|
"""
|
|
TODO: add docstring.
|
|
"""
|
|
return x_j * edge_attr.view(-1, 1)
|
|
|
|
def update(self, aggr_out, _):
|
|
"""
|
|
TODO: add docstring.
|
|
"""
|
|
return aggr_out
|
|
|
|
def aggregate(self, inputs, index, deg):
|
|
"""
|
|
TODO: add docstring.
|
|
"""
|
|
out = super().aggregate(inputs, index)
|
|
deg = deg + 1e-7
|
|
return out / deg.view(-1, 1)
|