def test_nvt_nose_hoover_jammed(self, dtype, sy_steps):
key = random.PRNGKey(0)
state = test_util.load_test_state('simulation_test_state.npy', dtype)
displacement_fn, shift_fn = space.periodic(state.box[0, 0])
E = energy.soft_sphere_pair(displacement_fn, state.species,
state.sigma)
invariant = partial(simulate.nvt_nose_hoover_invariant, E)
kT = 1e-3
init_fn, apply_fn = simulate.nvt_nose_hoover(E,
shift_fn,
1e-3,
kT=kT,
sy_steps=sy_steps)
apply_fn = jit(apply_fn)
state = init_fn(key, state.real_position)
E_initial = invariant(state, kT) * np.ones((DYNAMICS_STEPS, ))
def step_fn(i, state_and_energy):
state, energy = state_and_energy
state = apply_fn(state)
energy = ops.index_update(energy, i, invariant(state, kT))
return state, energy
Es = np.zeros((DYNAMICS_STEPS, ))
state, Es = lax.fori_loop(0, DYNAMICS_STEPS, step_fn, (state, Es))
tol = 1e-3 if dtype is f32 else 1e-7
self.assertEqual(state.position.dtype, dtype)
self.assertAllClose(Es, E_initial, rtol=tol, atol=tol)
def test_nvt_nose_hoover_ensemble(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
def invariant(T, state):
"""The conserved quantity for Nose-Hoover thermostat."""
accum = \
E(state.position) + quantity.kinetic_energy(state.velocity, state.mass)
DOF = spatial_dimension * PARTICLE_COUNT
accum = accum + (state.v_xi[0]) ** 2 * state.Q[0] * 0.5 + \
DOF * T * state.xi[0]
for xi, v_xi, Q in zip(state.xi[1:], state.v_xi[1:], state.Q[1:]):
accum = accum + v_xi**2 * Q * 0.5 + T * xi
return accum
for _ in range(STOCHASTIC_SAMPLES):
key, pos_key, center_key, vel_key, T_key, masses_key = \
random.split(key, 6)
R = random.normal(pos_key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
R0 = random.normal(center_key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
_, shift = space.free()
E = functools.partial(lambda R, R0, **kwargs: np.sum((R - R0)**2),
R0=R0)
T = random.uniform(T_key, (), minval=0.3, maxval=1.4, dtype=dtype)
mass = random.uniform(masses_key, (PARTICLE_COUNT, ),
minval=0.1,
maxval=10.0,
dtype=dtype)
init_fn, apply_fn = simulate.nvt_nose_hoover(E,
shift,
1e-3,
T,
tau=10)
apply_fn = jit(apply_fn)
state = init_fn(vel_key, R, mass=mass, T_initial=dtype(1.0))
initial = invariant(T, state)
for _ in range(DYNAMICS_STEPS):
state = apply_fn(state)
assert np.abs(
quantity.temperature(state.velocity, state.mass) - T) < 0.1
assert np.abs(invariant(T, state) - initial) < initial * 0.01
assert state.position.dtype == dtype
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)