shenoynikhil/annotated-dimenet

[Doubts] Radial Basis Layer

Closed this issue · 0 comments

Doubts on Radial Basis Layer

class RadialBasisLayer(torch.nn.Module):
    '''RadialBasisLayer'''
    def __init__(self, num_radial: int, cutoff: float = 5.0,
                 envelope_exponent: int = 5):
        super().__init__()
        # the c in the radial basis layer equation
        self.cutoff = cutoff
        # u(d) / envelope
        self.envelope = Envelope(envelope_exponent)
        
        # the different frequencies to be considered to generate orthogonal basis
        self.freq = torch.nn.Parameter(torch.Tensor(num_radial))
        
        # make sure we reset_parameters during __init__()
        self.reset_parameters()

    def reset_parameters(self):
        with torch.no_grad():
            torch.arange(1, self.freq.numel() + 1, out=self.freq).mul_(PI)
        self.freq.requires_grad_()

    def forward(self, dist: Tensor) -> Tensor:
        # compute d = (d/c)
        dist = (dist.unsqueeze(-1) / self.cutoff)
        # compute u(d/c) * sin(\frac{n\pi}{c} \times d)
        return self.envelope(dist) * (self.freq * dist).sin()
  • why not reset_parameters() with, torch.nn.Parameter(torch.arange(1, self.num_radial + 1).float())
  • why are we doing dist / self.cutoff in the forward loop and passing that to the envelope(), shouldn't it be envelope(d) *
  • I think the code equation looks like (based on forward()), but it is not what was shown in the paper,

$$ u(d / c) * sin(\frac{n\pi}{c}d) $$