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_cosine():
# cosine cutoff with radius 1.8
cutoff = CosineCutoff(cutoff=1.8)
# check cutoff radius
assert abs(1.8 - cutoff.cutoff) < 1.0e-12
# random tensor with elements in [0, 1)
torch.manual_seed(42)
dist = torch.rand((10, 5, 20), dtype=torch.float)
# check cutoff values
expt = 0.5 * (1.0 + torch.cos(dist * np.pi / 1.8))
assert torch.allclose(expt, cutoff(dist), atol=0.0, rtol=1.0e-7)
# compute expected values for 3.5 times distance
values = 0.5 * (1.0 + torch.cos(3.5 * dist * np.pi / 1.8))
values[3.5 * dist >= 1.8] = 0.0
assert torch.allclose(values, cutoff(3.5 * dist), atol=0.0, rtol=1.0e-7)
def test_cutoff_cosine_default():
# cosine cutoff with default radius
cutoff = CosineCutoff()
# check cutoff radius
assert abs(5.0 - cutoff.cutoff) < 1.0e-12
# random tensor with elements in [0, 1)
torch.manual_seed(42)
dist = torch.rand((1, 15), dtype=torch.float)
# check cutoff values
expt = 0.5 * (1.0 + torch.cos(dist * np.pi / 5.0))
assert torch.allclose(expt, cutoff(dist), atol=0.0, rtol=1.0e-7)
expt = 0.5 * (1.0 + torch.cos(2.0 * dist * np.pi / 5.0))
assert torch.allclose(expt, cutoff(2.0 * dist), atol=0.0, rtol=1.0e-7)
expt = 0.5 * (1.0 + torch.cos(4.0 * dist * np.pi / 5.0))
assert torch.allclose(expt, cutoff(4.0 * dist), atol=0.0, rtol=1.0e-7)
# compute cutoff values and expected values
comp = cutoff(5.5 * dist)
expt = 0.5 * (1.0 + torch.cos(5.5 * dist * np.pi / 5.0))
expt[5.5 * dist >= 5.0] = 0.0
assert torch.allclose(expt, comp, atol=0.0, rtol=1.0e-7)