def test_stress_non_minimized_periodic_general(self, dim, dtype, coords): key = random.PRNGKey(0) N = 64 box = quantity.box_size_at_number_density(N, 0.8, dim) displacement_fn, _ = space.periodic_general(box, coords == 'fractional') pos = random.uniform(key, (N, dim)) pos = pos if coords == 'fractional' else pos * box energy_fn = energy.soft_sphere_pair(displacement_fn) def exact_stress(R): dR = space.map_product(displacement_fn)(R, R) dr = space.distance(dR) g = jnp.vectorize(grad(energy.soft_sphere), signature='()->()') V = quantity.volume(dim, box) dUdr = 0.5 * g(dr)[:, :, None, None] dr = (dr + jnp.eye(N))[:, :, None, None] return -jnp.sum(dUdr * dR[:, :, None, :] * dR[:, :, :, None] / (V * dr), axis=(0, 1)) exact_stress = exact_stress(pos) ad_stress = quantity.stress(energy_fn, pos, box) tol = 1e-7 if dtype is f64 else 2e-5 self.assertAllClose(exact_stress, ad_stress, atol=tol, rtol=tol)
def test_pressure_non_minimized_free(self, dim, dtype): key = random.PRNGKey(0) N = 64 box = quantity.box_size_at_number_density(N, 0.8, dim) displacement_fn, _ = space.free() pos = random.uniform(key, (N, dim)) * box energy_fn = energy.soft_sphere_pair(displacement_fn) def exact_stress(R): dR = space.map_product(displacement_fn)(R, R) dr = space.distance(dR) g = jnp.vectorize(grad(energy.soft_sphere), signature='()->()') V = quantity.volume(dim, box) dUdr = 0.5 * g(dr)[:, :, None, None] dr = (dr + jnp.eye(N))[:, :, None, None] return jnp.sum(dUdr * dR[:, :, None, :] * dR[:, :, :, None] / (V * dr), axis=(0, 1)) exact_pressure = -1 / dim * jnp.trace(exact_stress(pos)) ad_pressure = quantity.pressure(energy_fn, pos, box) tol = 1e-7 if dtype is f64 else 2e-5 self.assertAllClose(exact_pressure, ad_pressure, atol=tol, rtol=tol)
def test_nvt_nose_hoover(self, spatial_dimension, dtype, sy_steps): key = random.PRNGKey(0) box_size = quantity.box_size_at_number_density(PARTICLE_COUNT, f32(1.2), spatial_dimension) displacement_fn, shift_fn = space.periodic(box_size) bonds_i = np.arange(PARTICLE_COUNT) bonds_j = np.roll(bonds_i, 1) bonds = np.stack([bonds_i, bonds_j]) E = energy.simple_spring_bond(displacement_fn, bonds) invariant = partial(simulate.nvt_nose_hoover_invariant, E) for _ in range(STOCHASTIC_SAMPLES): key, pos_key, vel_key, T_key, masses_key = random.split(key, 5) R = box_size * random.uniform(pos_key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype) T = random.uniform(T_key, (), minval=0.3, maxval=1.4, dtype=dtype) mass = 1 + random.uniform(masses_key, (PARTICLE_COUNT, ), dtype=dtype) init_fn, apply_fn = simulate.nvt_nose_hoover(E, shift_fn, 1e-3, T, sy_steps=sy_steps) apply_fn = jit(apply_fn) state = init_fn(vel_key, R, mass=mass) initial = invariant(state, T) for _ in range(DYNAMICS_STEPS): state = apply_fn(state) T_final = quantity.temperature(state.velocity, state.mass) assert np.abs(T_final - T) / T < 0.1 tol = 5e-4 if dtype is f32 else 1e-6 self.assertAllClose(invariant(state, T), initial, rtol=tol) self.assertEqual(state.position.dtype, dtype)
def make_periodic_general_test_system(N, dim, dtype, box_format): assert box_format in BOX_FORMATS box_size = quantity.box_size_at_number_density(N, 1.0, dim) box = dtype(box_size) if box_format == 'vector': box = jnp.array(jnp.ones(dim) * box_size, dtype) elif box_format == 'matrix': box = jnp.array(jnp.eye(dim) * box_size, dtype) d, s = space.periodic(jnp.diag(box) if box_format == 'matrix' else box) d_gf, s_gf = space.periodic_general(box) d_g, s_g = space.periodic_general(box, fractional_coordinates=False) key = random.PRNGKey(0) R_f = random.uniform(key, (N, dim), dtype=dtype) R = space.transform(box, R_f) E = jit(energy.soft_sphere_pair(d)) E_gf = jit(energy.soft_sphere_pair(d_gf)) E_g = jit(energy.soft_sphere_pair(d_g)) return R_f, R, box, (s, E), (s_gf, E_gf), (s_g, E_g)