def test_soft_sphere(self, spatial_dimension, alpha, dtype):
key = random.PRNGKey(0)
alpha = f32(alpha)
for _ in range(STOCHASTIC_SAMPLES):
key, split_sigma, split_epsilon = random.split(key, 3)
sigma = np.array(random.uniform(split_sigma, (1, ),
minval=0.0,
maxval=3.0)[0],
dtype=dtype)
epsilon = np.array(random.uniform(split_epsilon, (1, ),
minval=0.0,
maxval=4.0)[0],
dtype=dtype)
self.assertAllClose(
energy.soft_sphere(dtype(0), sigma, epsilon, alpha),
epsilon / alpha)
self.assertAllClose(
energy.soft_sphere(dtype(sigma), sigma, epsilon, alpha),
np.array(0.0, dtype=dtype))
self.assertAllClose(
grad(energy.soft_sphere)(dtype(2 * sigma), sigma, epsilon,
alpha), np.array(0.0, dtype=dtype))
if alpha > 2.0:
grad_energy = grad(energy.soft_sphere)
g = grad_energy(dtype(sigma), sigma, epsilon, alpha)
self.assertAllClose(g, np.array(0, dtype=dtype))
def test_swap_mc_jammed(self, dtype):
key = random.PRNGKey(0)
state = test_util.load_jammed_state('simulation_test_state.npy', dtype)
space_fn = space.periodic(state.box[0, 0])
displacement_fn, shift_fn = space_fn
sigma = np.diag(state.sigma)[state.species]
energy_fn = lambda dr, sigma: energy.soft_sphere(dr, sigma=sigma)
neighbor_fn = partition.neighbor_list(displacement_fn,
state.box[0, 0],
np.max(sigma) + 0.1,
dr_threshold=0.5)
kT = 1e-2
t_md = 0.1
N_swap = 10
init_fn, apply_fn = simulate.hybrid_swap_mc(space_fn,
energy_fn,
neighbor_fn,
1e-3,
kT,
t_md,
N_swap)
state = init_fn(key, state.real_position, sigma)
Ts = np.zeros((DYNAMICS_STEPS,))
def step_fn(i, state_and_temp):
state, temp = state_and_temp
state = apply_fn(state)
temp = temp.at[i].set(quantity.temperature(state.md.velocity))
return state, temp
state, Ts = lax.fori_loop(0, DYNAMICS_STEPS, step_fn, (state, Ts))
tol = 5e-4
self.assertAllClose(Ts[10:],
kT * np.ones((DYNAMICS_STEPS - 10)),
rtol=5e-1,
atol=5e-3)
self.assertAllClose(np.mean(Ts[10:]), kT, rtol=tol, atol=tol)
self.assertTrue(not np.all(state.sigma == sigma))
