def brownian(energy_or_force,
shift,
dt,
T_schedule,
quant=quantity.Energy,
gamma=0.1):
"""Simulation of Brownian dynamics.
This code simulates Brownian dynamics which are synonymous with the overdamped
regime of Langevin dynamics. However, in this case we don't need to take into
account velocity information and the dynamics simplify. Consequently, when
Brownian dynamics can be used they will be faster than Langevin. As in the
case of Langevin dynamics our implementation follows [1].
Args:
See nvt_langevin.
Returns:
See above.
[1] E. Carlon, M. Laleman, S. Nomidis. "Molecular Dynamics Simulation."
http://itf.fys.kuleuven.be/~enrico/Teaching/molecular_dynamics_2015.pdf
Accessed on 06/05/2019.
"""
force_fn = quantity.canonicalize_force(energy_or_force, quant)
dt, gamma = static_cast(dt, gamma)
T_schedule = interpolate.canonicalize(T_schedule)
def init_fn(key, R, mass=f32(1)):
mass = quantity.canonicalize_mass(mass)
return BrownianState(R, mass, key)
def apply_fn(state, t=f32(0), **kwargs):
R, mass, key = state
key, split = random.split(key)
F = force_fn(R, t=t, **kwargs)
xi = random.normal(split, R.shape, R.dtype)
nu = f32(1) / (mass * gamma)
dR = F * dt * nu + np.sqrt(f32(2) * T_schedule(t) * dt * nu) * xi
R = shift(R, dR, t=t, **kwargs)
return BrownianState(R, mass, key)
return init_fn, apply_fn
def nvt_langevin(energy_or_force,
shift,
dt,
T_schedule,
quant=quantity.Energy,
gamma=0.1):
"""Simulation in the NVT ensemble using the Langevin thermostat.
Samples from the canonical ensemble in which the number of particles (N),
the system volume (V), and the temperature (T) are held constant. Langevin
dynamics are stochastic and it is supposed that the system is interacting with
fictitious microscopic degrees of freedom. An example of this would be large
particles in a solvent such as water. Thus, Langevin dynamics are a stochastic
ODE described by a friction coefficient and noise of a given covariance.
Our implementation follows the excellent set of lecture notes by Carlon,
Laleman, and Nomidis [1].
Args:
energy_or_force: A function that produces either an energy or a force from
a set of particle positions specified as an ndarray of shape
[n, spatial_dimension].
shift_fn: A function that displaces positions, R, by an amount dR. Both R
and dR should be ndarrays of shape [n, spatial_dimension].
dt: Floating point number specifying the timescale (step size) of the
simulation.
T_schedule: Either a floating point number specifying a constant temperature
or a function specifying temperature as a function of time.
quant: Either a quantity.Energy or a quantity.Force specifying whether
energy_or_force is an energy or force respectively.
gamma: A float specifying the friction coefficient between the particles
and the solvent.
Returns:
See above.
[1] E. Carlon, M. Laleman, S. Nomidis. "Molecular Dynamics Simulation."
http://itf.fys.kuleuven.be/~enrico/Teaching/molecular_dynamics_2015.pdf
Accessed on 06/05/2019.
"""
force_fn = quantity.canonicalize_force(energy_or_force, quant)
dt_2 = dt / 2
dt2 = dt**2 / 2
dt32 = dt**(3.0 / 2.0) / 2
dt, dt_2, dt2, dt32, gamma = static_cast(dt, dt_2, dt2, dt32, gamma)
T_schedule = interpolate.canonicalize(T_schedule)
def init_fn(key, R, mass=f32(1), T_initial=f32(1)):
mass = quantity.canonicalize_mass(mass)
key, split = random.split(key)
V = np.sqrt(T_initial / mass) * random.normal(
split, R.shape, dtype=R.dtype)
V = V - np.mean(V, axis=0, keepdims=True)
return NVTLangevinState(R, V, force_fn(R, t=f32(0)), mass, key)
def apply_fn(state, t=f32(0), **kwargs):
R, V, F, mass, key = state
N, dim = R.shape
key, xi_key, theta_key = random.split(key, 3)
xi = random.normal(xi_key, (N, dim), dtype=R.dtype)
theta = random.normal(theta_key,
(N, dim), dtype=R.dtype) / np.sqrt(f32(3))
# NOTE(schsam): We really only need to recompute sigma if the temperature
# is nonconstant. @Optimization
# TODO(schsam): Check that this is really valid in the case that the masses
# are non identical for all particles.
sigma = np.sqrt(f32(2) * T_schedule(t) * gamma / mass)
C = dt2 * (F - gamma * V) + sigma * dt32 * (xi + theta)
R = shift(R, dt * V + F + C, t=t, **kwargs)
F_new = force_fn(R, t=t, **kwargs)
V = (f32(1) - dt * gamma) * V + dt_2 * (F_new + F)
V = V + sigma * np.sqrt(dt) * xi - gamma * C
return NVTLangevinState(R, V, F_new, mass, key)
return init_fn, apply_fn
def nvt_nose_hoover(energy_or_force,
shift_fn,
dt,
T_schedule,
quant=quantity.Energy,
chain_length=5,
tau=0.01):
"""Simulation in the NVT ensemble using a Nose Hoover Chain thermostat.
Samples from the canonical ensemble in which the number of particles (N),
the system volume (V), and the temperature (T) are held constant. We use a
Nose Hoover Chain thermostat described in [1, 2, 3]. We employ a similar
notation to [2] and the interested reader might want to look at that paper as
a reference.
Currently, the implementation only does a single timestep per Nose-Hoover
step. At some point we should support the multi-step case.
Args:
energy_or_force: A function that produces either an energy or a force from
a set of particle positions specified as an ndarray of shape
[n, spatial_dimension].
shift_fn: A function that displaces positions, R, by an amount dR. Both R
and dR should be ndarrays of shape [n, spatial_dimension].
dt: Floating point number specifying the timescale (step size) of the
simulation.
T_schedule: Either a floating point number specifying a constant temperature
or a function specifying temperature as a function of time.
quant: Either a quantity.Energy or a quantity.Force specifying whether
energy_or_force is an energy or force respectively.
chain_length: An integer specifying the length of the Nose-Hoover chain.
tau: A floating point timescale over which temperature equilibration occurs.
The performance of the Nose-Hoover chain thermostat is quite sensitive to
this choice.
Returns:
See above.
[1] Martyna, Glenn J., Michael L. Klein, and Mark Tuckerman.
"Nose-Hoover chains: The canonical ensemble via continuous dynamics."
The Journal of chemical physics 97, no. 4 (1992): 2635-2643.
[2] Martyna, Glenn, Mark Tuckerman, Douglas J. Tobias, and Michael L. Klein.
"Explicit reversible integrators for extended systems dynamics."
Molecular Physics 87. (1998) 1117-1157.
[3] Tuckerman, Mark E., Jose Alejandre, Roberto Lopez-Rendon,
Andrea L. Jochim, and Glenn J. Martyna.
"A Liouville-operator derived measure-preserving integrator for molecular
dynamics simulations in the isothermal-isobaric ensemble."
Journal of Physics A: Mathematical and General 39, no. 19 (2006): 5629.
"""
force = quantity.canonicalize_force(energy_or_force, quant)
dt_2 = dt / 2.0
dt_4 = dt_2 / 2.0
dt_8 = dt_4 / 2.0
dt, dt_2, dt_4, dt_8, tau = static_cast(dt, dt_2, dt_4, dt_8, tau)
T_schedule = interpolate.canonicalize(T_schedule)
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)
# 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)
def step_chain(KE, V, xi, v_xi, Q, DOF, T):
"""Applies a single update to the chain parameters and rescales velocity."""
M = chain_length - 1
# TODO(schsam): We can probably cache the G parameters from the previous
# update.
# TODO(schsam): It is also probably the case that we could do a better job
# of vectorizing this code.
G = (Q[M - 1] * v_xi[M - 1]**f32(2) - T) / Q[M]
v_xi = ops.index_add(v_xi, M, dt_4 * G)
for m in range(M - 1, 0, -1):
G = (Q[m - 1] * v_xi[m - 1]**f32(2) - T) / Q[m]
scale = np.exp(-dt_8 * v_xi[m + 1])
v_xi = ops.index_update(v_xi, m,
scale * (scale * v_xi[m] + dt_4 * G))
G = (f32(2.0) * KE - DOF * T) / Q[0]
scale = np.exp(-dt_8 * v_xi[1])
v_xi = ops.index_update(v_xi, 0, scale * (scale * v_xi[0] + dt_4 * G))
scale = np.exp(-dt_2 * v_xi[0])
KE = KE * scale**f32(2)
V = V * scale
xi = xi + dt_2 * v_xi
G = (f32(2) * KE - DOF * T) / Q[0]
for m in range(M):
scale = np.exp(-dt_8 * v_xi[m + 1])
v_xi = ops.index_update(v_xi, m,
scale * (scale * v_xi[m] + dt_4 * G))
G = (Q[m] * v_xi[m]**f32(2) - T) / Q[m + 1]
v_xi = ops.index_add(v_xi, M, dt_4 * G)
return KE, V, xi, v_xi
def apply_fun(state, t=0.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
# 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)
return init_fun, apply_fun
