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