def test_shape_cutoffs(distances):
cosine = CosineCutoff()
mollifier = MollifierCutoff()
hard = HardCutoff()
inputs = [distances]
out_shape = list(distances.shape)
assert_equal_shape(cosine, inputs, out_shape)
assert_equal_shape(mollifier, inputs, out_shape)
assert_equal_shape(hard, inputs, out_shape)
def test_cutoff_mollifier():
# mollifier cutoff with radius 2.3
cutoff = MollifierCutoff(cutoff=2.3)
# check cutoff radius
assert abs(2.3 - cutoff.cutoff) < 1.0e-12
# tensor of zeros
dist = torch.zeros((4, 1, 1))
assert torch.allclose(torch.ones(4, 1, 1),
cutoff(dist),
atol=0.0,
rtol=1.0e-7)
# random tensor with elements in [0, 1)
torch.manual_seed(42)
dist = torch.rand((1, 3, 9), dtype=torch.float)
# check cutoff values
expt = torch.exp(1.0 - 1.0 / (1.0 - (dist / 2.3)**2))
assert torch.allclose(expt, cutoff(dist), atol=0.0, rtol=1.0e-7)
# compute cutoff values and expected values
comp = cutoff(3.8 * dist)
expt = torch.exp(1.0 - 1.0 / (1.0 - (3.8 * dist / 2.3)**2))
expt[3.8 * dist >= 2.3] = 0.0
assert torch.allclose(expt, comp, atol=0.0, rtol=1.0e-7)
def test_cutoff_mollifier_default():
# mollifier cutoff with default radius
cutoff = MollifierCutoff()
# check cutoff radius
assert abs(5.0 - cutoff.cutoff) < 1.0e-12
# tensor of zeros
dist = torch.zeros((5, 2, 3))
assert torch.allclose(torch.ones(5, 2, 3),
cutoff(dist),
atol=0.0,
rtol=1.0e-7)
# random tensor with elements in [0, 1)
torch.manual_seed(42)
dist = torch.rand((20, 1), dtype=torch.float)
# check cutoff values
expt = torch.exp(1.0 - 1.0 / (1.0 - (dist / 5.0)**2))
assert torch.allclose(expt, cutoff(dist), atol=0.0, rtol=1.0e-7)
# compute cutoff values and expected values
comp = cutoff(6.0 * dist)
expt = torch.exp(1.0 - 1.0 / (1.0 - (6.0 * dist / 5.0)**2))
expt[6.0 * dist >= 5.0] = 0.0
assert torch.allclose(expt, comp, atol=0.0, rtol=1.0e-7)