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)