def init_fn(key, R, box, mass=f32(1.0), **kwargs):
N, dim = R.shape
_kT = kT if 'kT' not in kwargs else kwargs['kT']
mass = quantity.canonicalize_mass(mass)
V = jnp.sqrt(_kT / mass) * random.normal(key, R.shape, dtype=R.dtype)
V = V - jnp.mean(V * mass, axis=0, keepdims=True) / mass
KE = quantity.kinetic_energy(V, mass)
# The box position is defined via pos = (1 / d) log V / V_0.
zero = jnp.zeros((), dtype=R.dtype)
one = jnp.ones((), dtype=R.dtype)
box_position = zero
box_velocity = zero
box_mass = dim * (N + 1) * kT * barostat_kwargs['tau']**2 * one
KE_box = quantity.kinetic_energy(box_velocity, box_mass)
if jnp.isscalar(box) or box.ndim == 0:
# TODO(schsam): This is necessary because of JAX issue #5849.
box = jnp.eye(R.shape[-1]) * box
return NPTNoseHooverState(R, V, force_fn(R, box=box, **kwargs), mass,
box, box_position, box_velocity, box_mass,
barostat.initialize(1, KE_box, _kT),
thermostat.initialize(R.size, KE, _kT)) # pytype: disable=wrong-arg-count
def apply_fn(state, **kwargs):
S = state
_kT = kT if 'kT' not in kwargs else kwargs['kT']
bc = barostat.update_mass(S.barostat, _kT)
tc = thermostat.update_mass(S.thermostat, _kT)
S = update_box_mass(S, _kT)
V_b, bc = barostat.half_step(S.box_velocity, bc, _kT)
V, tc = thermostat.half_step(S.velocity, tc, _kT)
S = dataclasses.replace(S, velocity=V, box_velocity=V_b)
S = inner_step(S, **kwargs)
KE = quantity.kinetic_energy(S.velocity, S.mass)
tc = dataclasses.replace(tc, kinetic_energy=KE)
KE_box = quantity.kinetic_energy(S.box_velocity, S.box_mass)
bc = dataclasses.replace(bc, kinetic_energy=KE_box)
V, tc = thermostat.half_step(S.velocity, tc, _kT)
V_b, bc = barostat.half_step(S.box_velocity, bc, _kT)
S = dataclasses.replace(S,
thermostat=tc,
barostat=bc,
velocity=V,
box_velocity=V_b)
return S
def test_nve_ensemble_time_dependence(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
pos_key, center_key, vel_key, mass_key = random.split(key, 4)
R = random.normal(pos_key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
R0 = random.normal(center_key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
mass = random.uniform(mass_key, (PARTICLE_COUNT, ),
minval=0.1,
maxval=5.0,
dtype=dtype)
displacement, shift = space.free()
E = energy.soft_sphere_pair(displacement)
init_fn, apply_fn = simulate.nve(E, shift, 1e-3)
apply_fn = jit(apply_fn)
state = init_fn(vel_key, R, mass=mass)
E_T = lambda state: \
E(state.position) + quantity.kinetic_energy(state.velocity, state.mass)
E_initial = E_T(state)
for t in range(SHORT_DYNAMICS_STEPS):
state = apply_fn(state, t=t * 1e-3)
E_total = E_T(state)
assert np.abs(E_total - E_initial) < E_initial * 0.01
assert state.position.dtype == dtype
def nvt_nose_hoover_invariant(energy_fn: Callable[..., Array],
state: NVTNoseHooverState, kT: float,
**kwargs) -> float:
"""The conserved quantity for the NVT ensemble with a Nose-Hoover thermostat.
This function is normally used for debugging the Nose-Hoover thermostat.
Arguments:
energy_fn: The energy function of the Nose-Hoover system.
state: The current state of the system.
kT: The current goal temperature of the system.
Returns:
The Hamiltonian of the extended NVT dynamics.
"""
PE = energy_fn(state.position, **kwargs)
KE = quantity.kinetic_energy(state.velocity, state.mass)
DOF = state.position.size
E = PE + KE
c = state.chain
E += c.mass[0] * c.velocity[0]**2 / 2 + DOF * kT * c.position[0]
for r, v, m in zip(c.position[1:], c.velocity[1:], c.mass[1:]):
E += m * v**2 / 2 + kT * r
return E
def test_nve_ensemble(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
pos_key, center_key, vel_key, mass_key = random.split(key, 4)
R = random.normal(pos_key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
R0 = random.normal(center_key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
mass = random.uniform(mass_key, (PARTICLE_COUNT, ),
minval=0.1,
maxval=5.0,
dtype=dtype)
_, shift = space.free()
E = lambda R, **kwargs: np.sum((R - R0)**2)
init_fn, apply_fn = simulate.nve(E, shift, 1e-3)
apply_fn = jit(apply_fn)
state = init_fn(vel_key, R, mass=mass)
E_T = lambda state: \
E(state.position) + quantity.kinetic_energy(state.velocity, state.mass)
E_initial = E_T(state)
for _ in range(DYNAMICS_STEPS):
state = apply_fn(state)
E_total = E_T(state)
assert np.abs(E_total - E_initial) < E_initial * 0.01
assert state.position.dtype == dtype
def test_nve_jammed(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
state = test_util.load_jammed_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)
init_fn, apply_fn = simulate.nve(E, shift_fn, 1e-3)
apply_fn = jit(apply_fn)
state = init_fn(key, state.real_position, kT=1e-3)
E_T = lambda state: \
E(state.position) + quantity.kinetic_energy(state.velocity, state.mass)
E_initial = E_T(state) * np.ones((DYNAMICS_STEPS, ))
def step_fn(i, state_and_energy):
state, energy = state_and_energy
state = apply_fn(state)
energy = energy.at[i].set(E_T(state))
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_nve_jammed_periodic_general(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)
init_fn, apply_fn = simulate.nve(E, shift_fn, 1e-3)
apply_fn = jit(apply_fn)
state = init_fn(key, getattr(state, coords + '_position'), kT=1e-3)
E_T = lambda state: \
E(state.position) + quantity.kinetic_energy(state.velocity, state.mass)
E_initial = E_T(state) * 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, E_T(state))
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 nose_hoover_invariant(energy_fn: Callable[..., Array],
state: NVTNoseHooverState, kT: float,
**kwargs) -> float:
"""The conserved quantity for the Nose-Hoover thermostat.
This function is normally used for debugging the Nose-Hoover thermostat.
Arguments:
energy_fn: The energy function of the Nose-Hoover system.
state: The current state of the system.
kT: The current goal temperature of the system.
Returns:
The Hamiltonian of the extended NVT dynamics.
"""
PE = energy_fn(state.position, **kwargs)
KE = quantity.kinetic_energy(state.velocity, state.mass)
DOF = state.position.size
E = PE + KE
E += state.v_xi[0]**2 * state.Q[0] * 0.5 + DOF * kT * state.xi[0]
for xi, v_xi, Q in zip(state.xi[1:], state.v_xi[1:], state.Q[1:]):
E += v_xi**2 * Q * 0.5 + kT * xi
return E
def step_fn(i, state_energy_pressure):
state, energy, pressure = state_energy_pressure
state = apply_fn(state)
energy = energy.at[i].set(invariant(state, P, kT))
box = simulate.npt_box(state)
KE = quantity.kinetic_energy(state.velocity, state.mass)
p = pressure_fn(state.position, box, KE)
pressure = pressure.at[i].set(p)
return state, energy, pressure
def step_fn(i, state_energy_pressure):
state, energy, pressure = state_energy_pressure
state = apply_fn(state)
energy = ops.index_update(energy, i, invariant(state, P, kT))
box = simulate.npt_box(state)
KE = quantity.kinetic_energy(state.velocity, state.mass)
p = pressure_fn(state.position, box, KE)
pressure = ops.index_update(pressure, i, p)
return state, energy, pressure
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
def init_fn(key, R, mass=f32(1.0), **kwargs):
_kT = kT if 'kT' not in kwargs else kwargs['kT']
mass = quantity.canonicalize_mass(mass)
V = jnp.sqrt(_kT / mass) * random.normal(key, R.shape, dtype=R.dtype)
V = V - jnp.mean(V * mass, axis=0, keepdims=True) / mass
KE = quantity.kinetic_energy(V, mass)
return NVTNoseHooverState(R, V, force_fn(R, **kwargs), mass,
chain_fns.initialize(R.size, KE, _kT)) # pytype: disable=wrong-arg-count
def init_fun(key, R, mass=f32(1.0), T_initial=f32(1.0)):
mass = quantity.canonicalize_mass(mass)
V = np.sqrt(T_initial / mass) * random.normal(
key, R.shape, dtype=R.dtype)
V = V - np.mean(V, axis=0, keepdims=True)
KE = quantity.kinetic_energy(V, mass)
# Nose-Hoover parameters.
xi = np.zeros(chain_length, R.dtype)
v_xi = np.zeros(chain_length, R.dtype)
DOF, = static_cast(R.shape[0] * R.shape[1])
Q = T_initial * tau**f32(2) * np.ones(chain_length, dtype=R.dtype)
Q = ops.index_update(Q, 0, Q[0] * DOF)
return NVTNoseHooverState(R, V, mass, KE, xi, v_xi, Q)
def apply_fn(state, **kwargs):
_kT = kT if 'kT' not in kwargs else kwargs['kT']
chain = state.chain
chain = chain_fns.update_mass(chain, _kT)
v, chain = chain_fns.half_step(state.velocity, chain, _kT)
state = dataclasses.replace(state, velocity=v)
state = velocity_verlet(force_fn, shift_fn, dt, state, **kwargs)
KE = quantity.kinetic_energy(state.velocity, state.mass)
chain = dataclasses.replace(chain, kinetic_energy=KE)
v, chain = chain_fns.half_step(state.velocity, chain, _kT)
state = dataclasses.replace(state, velocity=v, chain=chain)
return state
def init_fun(key, R, mass=f32(1.0), T_initial=None):
if T_initial is None:
T_initial = T_schedule(0.0)
mass = quantity.canonicalize_mass(mass)
V = np.sqrt(T_initial / mass) * random.normal(
key, R.shape, dtype=R.dtype)
V = V - np.mean(V, axis=0, keepdims=True)
KE = quantity.kinetic_energy(V, mass)
# Nose-Hoover parameters.
xi = np.zeros(chain_length, R.dtype)
v_xi = np.zeros(chain_length, R.dtype)
# TODO(schsam): Really, it seems like Q should be set by the goal
# temperature rather than the initial temperature.
DOF, = static_cast(R.shape[0] * R.shape[1])
Q = T_initial * tau**f32(2) * np.ones(chain_length, dtype=R.dtype)
Q = ops.index_update(Q, 0, Q[0] * DOF)
return NVTNoseHooverState(R, V, mass, KE, xi, v_xi, Q) # pytype: disable=wrong-arg-count
def init_fn(key, R, mass=f32(1.0), **kwargs):
_kT = kT if 'kT' not in kwargs else kwargs['kT']
mass = quantity.canonicalize_mass(mass)
V = np.sqrt(_kT / mass) * random.normal(key, R.shape, dtype=R.dtype)
V = V - np.mean(V, axis=0, keepdims=True)
KE = quantity.kinetic_energy(V, mass)
# Nose-Hoover parameters.
xi = np.zeros(chain_length, R.dtype)
v_xi = np.zeros(chain_length, R.dtype)
# TODO(schsam): Really, it seems like Q should be set by the goal
# temperature rather than the initial temperature.
DOF = f32(R.shape[0] * R.shape[1])
Q = _kT * tau**f32(2) * np.ones(chain_length, dtype=R.dtype)
Q = ops.index_update(Q, 0, Q[0] * DOF)
F = force_fn(R, **kwargs)
return NVTNoseHooverState(R, V, F, mass, KE, xi, v_xi, Q) # pytype: disable=wrong-arg-count
def apply_fun(state, t=f32(0), **kwargs):
T = T_schedule(t)
R, V, mass, KE, xi, v_xi, Q = state
DOF, = static_cast(R.shape[0] * R.shape[1])
Q = T * tau**f32(2) * np.ones(chain_length, dtype=R.dtype)
Q = ops.index_update(Q, 0, Q[0] * DOF)
KE, V, xi, v_xi = step_chain(KE, V, xi, v_xi, Q, DOF, T)
R = shift_fn(R, V * dt_2, t=t, **kwargs)
F = force(R, t=t, **kwargs)
V = V + dt * F / mass
V = V - np.mean(V, axis=0, keepdims=True)
KE = quantity.kinetic_energy(V, mass)
R = shift_fn(R, V * dt_2, t=t, **kwargs)
KE, V, xi, v_xi = step_chain(KE, V, xi, v_xi, Q, DOF, T)
return NVTNoseHooverState(R, V, mass, KE, xi, v_xi, Q)
def apply_fn(state, **kwargs):
_kT = kT if 'kT' not in kwargs else kwargs['kT']
R, V, F, mass, KE, xi, v_xi, Q = dataclasses.astuple(state)
DOF = R.size
Q = _kT * tau**f32(2) * np.ones(chain_length, dtype=R.dtype)
Q = ops.index_update(Q, 0, Q[0] * DOF)
KE, V, xi, v_xi, *_ = half_step_chain_fn(KE, V, xi, v_xi, Q, DOF, _kT)
R = shift_fn(R, V * dt + F * dt**2 / (2 * mass), **kwargs)
F_new = force_fn(R, **kwargs)
V = V + dt_2 * (F_new + F) / mass
V = V - np.mean(V, axis=0, keepdims=True)
KE = quantity.kinetic_energy(V, mass)
KE, V, xi, v_xi, *_ = half_step_chain_fn(KE, V, xi, v_xi, Q, DOF, _kT)
return NVTNoseHooverState(R, V, F_new, mass, KE, xi, v_xi, Q)
def apply_fun(state, t=f32(0), **kwargs):
T = T_schedule(t)
R, V, mass, KE, xi, v_xi, Q = dataclasses.astuple(state)
DOF, = static_cast(R.shape[0] * R.shape[1])
Q = T * tau**f32(2) * np.ones(chain_length, dtype=R.dtype)
Q = ops.index_update(Q, 0, Q[0] * DOF)
KE, V, xi, v_xi = step_chain(KE, V, xi, v_xi, Q, DOF, T)
R = shift_fn(R, V * dt_2, t=t, **kwargs)
F = force(R, t=t, **kwargs)
V = V + dt * F / mass
# NOTE(schsam): Do we need to mean subtraction here?
V = V - np.mean(V, axis=0, keepdims=True)
KE = quantity.kinetic_energy(V, mass)
R = shift_fn(R, V * dt_2, t=t, **kwargs)
KE, V, xi, v_xi = step_chain(KE, V, xi, v_xi, Q, DOF, T)
return NVTNoseHooverState(R, V, mass, KE, xi, v_xi, Q) # pytype: disable=wrong-arg-count
def npt_nose_hoover_invariant(energy_fn: Callable[..., Array],
state: NPTNoseHooverState, pressure: float,
kT: float, **kwargs) -> float:
"""The conserved quantity for the NPT ensemble with a Nose-Hoover thermostat.
This function is normally used for debugging the NPT simulation.
Arguments:
energy_fn: The energy function of the system.
state: The current state of the system.
pressure: The current goal pressure of the system.
kT: The current goal temperature of the system.
Returns:
The Hamiltonian of the extended NPT dynamics.
"""
volume, box_fn = _npt_box_info(state)
PE = energy_fn(state.position, box=box_fn(volume), **kwargs)
KE = quantity.kinetic_energy(state.velocity, state.mass)
DOF = state.position.size
E = PE + KE
c = state.thermostat
E += c.mass[0] * c.velocity[0]**2 / 2 + DOF * kT * c.position[0]
for r, v, m in zip(c.position[1:], c.velocity[1:], c.mass[1:]):
E += m * v**2 / 2 + kT * r
c = state.barostat
for r, v, m in zip(c.position, c.velocity, c.mass):
E += m * v**2 / 2 + kT * r
E += pressure * volume
E += state.box_mass * state.box_velocity**2 / 2
return E
def do_fn(theta):
key = random.PRNGKey(0)
V = random.normal(key, (PARTICLE_COUNT, spatial_dimension),
dtype=f32)
return quantity.kinetic_energy(theta * V)
init, apply = simulate.nve(energy_fn, shift_fn, args.time / args.steps)
apply = jit(apply)
state = init(key, R, velocity_scale=0.0)
PE = []
KE = []
print_every = args.log
old_time = time.perf_counter()
print('Step\tKE\tPE\tTotal Energy\ttime/step')
print('----------------------------------------')
for i in range(args.steps // print_every):
state = lax.fori_loop(0, print_every, lambda _, state: apply(state), state)
PE += [energy_fn(state.position)]
KE += [quantity.kinetic_energy(state.velocity)]
new_time = time.perf_counter()
print('{}\t{:.2f}\t{:.2f}\t{:.3f}\t{:.2f}'.format(
i * print_every, KE[-1], PE[-1], KE[-1] + PE[-1],
(new_time - old_time) / print_every / 10.0))
old_time = new_time
PE = np.array(PE)
KE = np.array(KE)
R = state.position
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns