def test_isotropic_cutoff(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
for _ in range(STOCHASTIC_SAMPLES):
key, split_rs, split_rl, split_sigma, split_epsilon = random.split(
key, 5)
sigma = f32(
random.uniform(split_sigma, (1, ), minval=0.5, maxval=3.0)[0])
epsilon = f32(
random.uniform(split_epsilon, (1, ), minval=0.0,
maxval=4.0)[0])
r_small = random.uniform(split_rs, (10, ),
minval=0.0,
maxval=2.0 * sigma,
dtype=dtype)
r_large = random.uniform(split_rl, (10, ),
minval=2.5 * sigma,
maxval=3.0 * sigma,
dtype=dtype)
r_onset = f32(2.0 * sigma)
r_cutoff = f32(2.5 * sigma)
E = energy.multiplicative_isotropic_cutoff(energy.lennard_jones,
r_onset, r_cutoff)
self.assertAllClose(E(r_small, sigma, epsilon),
energy.lennard_jones(r_small, sigma, epsilon))
self.assertAllClose(E(r_large, sigma, epsilon),
np.zeros_like(r_large, dtype=dtype))
def test_lennard_jones(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
for _ in range(STOCHASTIC_SAMPLES):
key, split_sigma, split_epsilon = random.split(key, 3)
sigma = dtype(
random.uniform(split_sigma, (1, ), minval=0.5, maxval=3.0)[0])
epsilon = dtype(
random.uniform(split_epsilon, (1, ), minval=0.0,
maxval=4.0)[0])
dr = dtype(sigma * 2**(1.0 / 6.0))
self.assertAllClose(energy.lennard_jones(dr, sigma, epsilon),
np.array(-epsilon, dtype=dtype))
g = grad(energy.lennard_jones)(dr, sigma, epsilon)
self.assertAllClose(g, np.array(0, dtype=dtype))
def test_stress_lammps_periodic_general(self, dim, dtype):
key = random.PRNGKey(0)
N = 64
(box, R, V), (E, C) = test_util.load_lammps_stress_data(dtype)
displacement_fn, _ = space.periodic_general(box)
energy_fn = smap.pair(
lambda dr, **kwargs: jnp.where(dr < f32(2.5),
energy.lennard_jones(dr), f32(0.0)),
space.canonicalize_displacement_or_metric(displacement_fn))
ad_stress = quantity.stress(energy_fn, R, box, velocity=V)
tol = 5e-5
self.assertAllClose(energy_fn(R) / len(R), E, atol=tol, rtol=tol)
self.assertAllClose(C, ad_stress, atol=tol, rtol=tol)
