def test_nvt_langevin(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
for _ in range(STOCHASTIC_SAMPLES):
key, R_key, R0_key, T_key, masses_key = random.split(key, 5)
R = random.normal(
R_key, (LANGEVIN_PARTICLE_COUNT, spatial_dimension), dtype=dtype)
R0 = random.normal(
R0_key, (LANGEVIN_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, (LANGEVIN_PARTICLE_COUNT,), minval=0.1, maxval=10.0, dtype=dtype)
init_fn, apply_fn = simulate.nvt_langevin(E, shift, f32(1e-2), T, gamma=f32(0.3))
apply_fn = jit(apply_fn)
state = init_fn(key, R, mass=mass, T_initial=dtype(1.0))
T_list = []
for step in range(LANGEVIN_DYNAMICS_STEPS):
state = apply_fn(state)
if step > 4000 and step % 100 == 0:
T_list += [quantity.temperature(state.velocity, state.mass)]
T_emp = np.mean(np.array(T_list))
assert np.abs(T_emp - T) < 0.1
assert state.position.dtype == dtype
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)
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
