def test_pressure_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_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_pressure_jammed(self, dtype, coords):
key = random.PRNGKey(0)
state = test_util.load_test_state('simulation_test_state.npy', dtype)
displacement_fn, shift_fn = space.periodic_general(state.box,
coords == 'fractional')
E = energy.soft_sphere_pair(displacement_fn, state.species, state.sigma)
pos = getattr(state, coords + '_position')
self.assertAllClose(quantity.pressure(E, pos, state.box), state.pressure)
def test_pressure_jammed_periodic(self, dtype):
key = random.PRNGKey(0)
state = test_util.load_jammed_state('simulation_test_state.npy', dtype)
displacement_fn, shift_fn = space.periodic(jnp.diag(state.box))
E = energy.soft_sphere_pair(displacement_fn, state.species, state.sigma)
pos = state.real_position
tol = 1e-7 if dtype is f64 else 2e-5
self.assertAllClose(quantity.pressure(E, pos, state.box), state.pressure,
atol=tol, rtol=tol)