File size: 960 Bytes
3e02ab8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | # This file is a part of the `nequip` package. Please see LICENSE and README at the root for information on using it.
import torch
class PolynomialCutoff(torch.nn.Module):
def __init__(self, p: float = 6):
r"""Polynomial cutoff, as proposed in DimeNet: https://arxiv.org/abs/2003.03123
Args:
r_max (float): cutoff radius
p (int) : power used in envelope function
"""
super().__init__()
assert p >= 2.0
self.p = float(p)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Evaluate cutoff function.
Args:
x (torch.Tensor): input distance
"""
out = 1.0
out = out - (((self.p + 1.0) * (self.p + 2.0) / 2.0) * torch.pow(x, self.p))
out = out + (self.p * (self.p + 2.0) * torch.pow(x, self.p + 1.0))
out = out - ((self.p * (self.p + 1.0) / 2) * torch.pow(x, self.p + 2.0))
return out * (x < 1.0)
|