dmlc--dgl
Chang Liu
31 行
982 B
Python
31 行
982 B
Python
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from modules.envelope import Envelope
|
|
|
|
class BesselBasisLayer(nn.Module):
|
|
def __init__(self,
|
|
num_radial,
|
|
cutoff,
|
|
envelope_exponent=5):
|
|
super(BesselBasisLayer, self).__init__()
|
|
|
|
self.cutoff = cutoff
|
|
self.envelope = Envelope(envelope_exponent)
|
|
self.frequencies = nn.Parameter(torch.Tensor(num_radial))
|
|
self.reset_params()
|
|
|
|
def reset_params(self):
|
|
with torch.no_grad():
|
|
torch.arange(1, self.frequencies.numel() + 1, out=self.frequencies).mul_(np.pi)
|
|
self.frequencies.requires_grad_()
|
|
|
|
def forward(self, g):
|
|
d_scaled = g.edata['d'] / self.cutoff
|
|
# Necessary for proper broadcasting behaviour
|
|
d_scaled = torch.unsqueeze(d_scaled, -1)
|
|
d_cutoff = self.envelope(d_scaled)
|
|
g.edata['rbf'] = d_cutoff * torch.sin(self.frequencies * d_scaled)
|
|
return g
|