def gradient_descent(energy_or_force: Callable[..., Array], shift_fn: ShiftFn,
step_size: float) -> Minimizer[Array]:
"""Defines gradient descent minimization.
This is the simplest optimization strategy that moves particles down their
gradient to the nearest minimum. Generally, gradient descent is slower than
other methods and is included mostly for its simplicity.
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].
step_size: A floating point specifying the size of each step.
quant: Either a quantity.Energy or a quantity.Force specifying whether
energy_or_force is an energy or force respectively.
Returns:
See above.
"""
force = quantity.canonicalize_force(energy_or_force)
def init_fn(R: Array, **unused_kwargs) -> Array:
return R
def apply_fn(R: Array, **kwargs) -> Array:
R = shift_fn(R, step_size * force(R, **kwargs), **kwargs)
return R
return init_fn, apply_fn
def nve(energy_or_force, shift_fn, dt, quant=quantity.Energy):
"""Simulates a system in the NVE ensemble.
Samples from the microcanonical ensemble in which the number of particles (N),
the system volume (V), and the energy (E) are held constant. We use a standard
velocity verlet integration scheme.
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.
quant: Either a quantity.Energy or a quantity.Force specifying whether
energy_or_force is an energy or force respectively.
Returns:
See above.
"""
force = quantity.canonicalize_force(energy_or_force, quant)
dt, = static_cast(dt)
dt_2, = static_cast(0.5 * dt ** 2)
def init_fun(key, R, velocity_scale=f32(1.0), mass=f32(1.0)):
V = np.sqrt(velocity_scale) * random.normal(key, R.shape, dtype=R.dtype)
mass = quantity.canonicalize_mass(mass)
return NVEState(R, V, force(R) / mass, mass)
def apply_fun(state, t=None, **kwargs):
R, V, A, mass = state
R = shift_fn(R, V * dt + A * dt_2, t=t, **kwargs)
A_prime = force(R, t=t, **kwargs) / mass
V = V + f32(0.5) * (A + A_prime) * dt
return NVEState(R, V, A_prime, mass)
return init_fun, apply_fun
def nve(energy_or_force_fn: Callable[..., Array], shift_fn: ShiftFn,
dt: float) -> Simulator:
"""Simulates a system in the NVE ensemble.
Samples from the microcanonical ensemble in which the number of particles
(N), the system volume (V), and the energy (E) are held constant. We use a
standard velocity verlet integration scheme.
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.
quant: Either a quantity.Energy or a quantity.Force specifying whether
energy_or_force is an energy or force respectively.
Returns:
See above.
"""
force_fn = quantity.canonicalize_force(energy_or_force_fn)
def init_fn(key, R, kT, mass=f32(1.0), **kwargs):
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
return NVEState(R, V, force_fn(R, **kwargs), mass) # pytype: disable=wrong-arg-count
def step_fn(state, **kwargs):
return velocity_verlet(force_fn, shift_fn, dt, state, **kwargs)
return init_fn, step_fn
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 brownian(energy_or_force: Callable[..., Array],
shift: ShiftFn,
dt: float,
kT: float,
gamma: float = 0.1) -> Simulator:
"""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:
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.
kT: Floating point number specifying the temperature inunits of Boltzmann
constant. To update the temperature dynamically during a simulation one
should pass `kT` as a keyword argument to the step function.
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)
dt, gamma = static_cast(dt, gamma)
def init_fn(key, R, mass=f32(1)):
mass = quantity.canonicalize_mass(mass)
return BrownianState(R, mass, key) # pytype: disable=wrong-arg-count
def apply_fn(state, **kwargs):
_kT = kT if 'kT' not in kwargs else kwargs['kT']
R, mass, key = dataclasses.astuple(state)
key, split = random.split(key)
F = force_fn(R, **kwargs)
xi = random.normal(split, R.shape, R.dtype)
nu = f32(1) / (mass * gamma)
dR = F * dt * nu + np.sqrt(f32(2) * _kT * dt * nu) * xi
R = shift(R, dR, **kwargs)
return BrownianState(R, mass, key) # pytype: disable=wrong-arg-count
return init_fn, apply_fn
def nvt_langevin(energy_or_force: Callable[..., Array],
shift: ShiftFn,
dt: float,
kT: float,
gamma: float = 0.1) -> Simulator:
"""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.
kT: Floating point number specifying the temperature inunits of Boltzmann
constant. To update the temperature dynamically during a simulation one
should pass `kT` as a keyword argument to the step function.
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)
dt_2 = f32(dt / 2)
dt2 = f32(dt**2 / 2)
dt32 = f32(dt**(3.0 / 2.0) / 2)
kT = f32(kT)
gamma = f32(gamma)
def init_fn(key, R, mass=f32(1), **kwargs):
_kT = kT if 'kT' not in kwargs else kwargs['kT']
mass = quantity.canonicalize_mass(mass)
key, split = random.split(key)
V = np.sqrt(_kT / 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, **kwargs), mass, key) # pytype: disable=wrong-arg-count
def apply_fn(state, **kwargs):
R, V, F, mass, key = dataclasses.astuple(state)
_kT = kT if 'kT' not in kwargs else kwargs['kT']
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) * _kT * gamma / mass)
C = dt2 * (F - gamma * V) + sigma * dt32 * (xi + theta)
R = shift(R, dt * V + C, **kwargs)
F_new = force_fn(R, **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) # pytype: disable=wrong-arg-count
return init_fn, apply_fn
def nvt_nose_hoover(energy_or_force: Callable[..., Array],
shift_fn: ShiftFn,
dt: float,
kT: float,
chain_length: int = 5,
chain_steps: int = 2,
sy_steps: int = 3,
tau: float = None) -> Simulator:
"""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 (NHC) 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.
As described in [3], the NHC evolves on a faster timescale than the rest of
the simulation. Therefore, it often desirable to integrate the chain over
several substeps for each step of MD. To do this we follow the Suzuki-Yoshida
scheme. Specifically, we subdivide our chain simulation into $n_c$ substeps.
These substeps are further subdivided into $n_sy$ steps. Each $n_sy$ step has
length $\delta_i = \Delta t w_i / n_c$ where $w_i$ are constants such that
$\sum_i w_i = 1$. See the table of Suzuki_Yoshida weights above for specific
values.
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.
kT: Floating point number specifying the temperature inunits of Boltzmann
constant. To update the temperature dynamically during a simulation one
should pass `kT` as a keyword argument to the step function.
chain_length: An integer specifying the number of particles in
the Nose-Hoover chain.
chain_steps: An integer specifying the number, $n_c$, of outer substeps.
sy_steps: An integer specifying the number of Suzuki-Yoshida steps. This
must be either 1, 3, 5, or 7.
tau: A floating point timescale over which temperature equilibration occurs.
Measured in units of dt. The performance of the Nose-Hoover chain
thermostat can be 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_fn = quantity.canonicalize_force(energy_or_force)
dt = f32(dt)
if tau is None:
tau = dt * 100
tau = f32(tau)
dt_2 = dt / f32(2.0)
kT = f32(kT)
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 substep_chain_fn(delta, KE, V, xi, v_xi, Q, DOF, T):
"""Applies a single update to the chain parameters and rescales velocity."""
delta_2 = delta / f32(2.0)
delta_4 = delta_2 / f32(2.0)
delta_8 = delta_4 / f32(2.0)
M = chain_length - 1
G = (Q[M - 1] * v_xi[M - 1]**f32(2) - T) / Q[M]
v_xi = ops.index_add(v_xi, M, delta_4 * G)
def backward_loop_fn(v_xi_new, m):
G = (Q[m - 1] * v_xi[m - 1]**2 - T) / Q[m]
scale = np.exp(-delta_8 * v_xi_new)
v_xi_new = scale * (scale * v_xi[m] + delta_4 * G)
return v_xi_new, v_xi_new
idx = np.arange(M - 1, 0, -1)
_, v_xi_update = lax.scan(backward_loop_fn, v_xi[M], idx, unroll=2)
v_xi = ops.index_update(v_xi, idx, v_xi_update)
G = (f32(2.0) * KE - DOF * T) / Q[0]
scale = np.exp(-delta_8 * v_xi[1])
v_xi = ops.index_update(v_xi, 0,
scale * (scale * v_xi[0] + delta_4 * G))
scale = np.exp(-delta_2 * v_xi[0])
KE = KE * scale**f32(2)
V = V * scale
xi = xi + delta_2 * v_xi
G = (f32(2) * KE - DOF * T) / Q[0]
def forward_loop_fn(G, m):
scale = np.exp(-delta_8 * v_xi[m + 1])
v_xi_update = scale * (scale * v_xi[m] + delta_4 * G)
G = (Q[m] * v_xi_update**f32(2) - T) / Q[m + 1]
return G, v_xi_update
idx = np.arange(M)
G, v_xi_update = lax.scan(forward_loop_fn, G, idx, unroll=2)
v_xi = ops.index_update(v_xi, idx, v_xi_update)
v_xi = ops.index_add(v_xi, M, delta_4 * G)
return KE, V, xi, v_xi, Q, DOF, T
def half_step_chain_fn(*chain_state):
if chain_steps == 1 and sy_steps == 1:
return substep_chain_fn(dt, *chain_state)
delta = dt / chain_steps
ws = np.array(SUZUKI_YOSHIDA_WEIGHTS[sy_steps],
dtype=chain_state[1].dtype)
return lax.scan(
lambda chain_state, i:
(substep_chain_fn(delta * ws[i % sy_steps], *chain_state), 0),
chain_state, np.arange(chain_steps * sy_steps))[0]
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)
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
def fire_descent(energy_or_force: Callable[..., Array],
shift_fn: ShiftFn,
dt_start: float = 0.1,
dt_max: float = 0.4,
n_min: float = 5,
f_inc: float = 1.1,
f_dec: float = 0.5,
alpha_start: float = 0.1,
f_alpha: float = 0.99) -> Minimizer[FireDescentState]:
"""Defines FIRE minimization.
This code implements the "Fast Inertial Relaxation Engine" from [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].
quant: Either a quantity.Energy or a quantity.Force specifying whether
energy_or_force is an energy or force respectively.
dt_start: The initial step size during minimization as a float.
dt_max: The maximum step size during minimization as a float.
n_min: An integer specifying the minimum number of steps moving in the
correct direction before dt and f_alpha should be updated.
f_inc: A float specifying the fractional rate by which the step size
should be increased.
f_dec: A float specifying the fractional rate by which the step size
should be decreased.
alpha_start: A float specifying the initial momentum.
f_alpha: A float specifying the fractional change in momentum.
Returns:
See above.
[1] Bitzek, Erik, Pekka Koskinen, Franz Gahler, Michael Moseler,
and Peter Gumbsch. "Structural relaxation made simple."
Physical review letters 97, no. 17 (2006): 170201.
"""
dt_start, dt_max, n_min, f_inc, f_dec, alpha_start, f_alpha = util.static_cast(
dt_start, dt_max, n_min, f_inc, f_dec, alpha_start, f_alpha)
force = quantity.canonicalize_force(energy_or_force)
def init_fn(R: Array, **kwargs) -> FireDescentState:
V = jnp.zeros_like(R)
n_pos = jnp.zeros((), jnp.int32)
F = force(R, **kwargs)
return FireDescentState(R, V, F, dt_start, alpha_start, n_pos) # pytype: disable=wrong-arg-count
def apply_fn(state: FireDescentState, **kwargs) -> FireDescentState:
R, V, F_old, dt, alpha, n_pos = dataclasses.astuple(state)
R = shift_fn(R, dt * V + dt**f32(2) * F_old, **kwargs)
F = force(R, **kwargs)
V = V + dt * f32(0.5) * (F_old + F)
# NOTE(schsam): This will be wrong if F_norm ~< 1e-8.
# TODO(schsam): We should check for forces below 1e-6. @ErrorChecking
F_norm = jnp.sqrt(jnp.sum(F**f32(2)) + f32(1e-6))
V_norm = jnp.sqrt(jnp.sum(V**f32(2)))
P = jnp.array(jnp.dot(jnp.reshape(F, (-1)), jnp.reshape(V, (-1))))
V = V + alpha * (F * V_norm / F_norm - V)
# NOTE(schsam): Can we clean this up at all?
n_pos = jnp.where(P >= 0, n_pos + 1, 0)
dt_choice = jnp.array([dt * f_inc, dt_max])
dt = jnp.where(P > 0, jnp.where(n_pos > n_min, jnp.min(dt_choice), dt),
dt)
dt = jnp.where(P < 0, dt * f_dec, dt)
alpha = jnp.where(P > 0,
jnp.where(n_pos > n_min, alpha * f_alpha, alpha),
alpha)
alpha = jnp.where(P < 0, alpha_start, alpha)
V = (P < 0) * jnp.zeros_like(V) + (P >= 0) * V
return FireDescentState(R, V, F, dt, alpha, n_pos) # pytype: disable=wrong-arg-count
return init_fn, apply_fn
def fire_descent(energy_or_force,
shift_fn,
quant=quantity.Energy,
dt_start=0.1,
dt_max=0.4,
n_min=5,
f_inc=1.1,
f_dec=0.5,
alpha_start=0.1,
f_alpha=0.99):
"""Defines FIRE minimization.
This code implements the "Fast Inertial Relaxation Engine" from [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].
quant: Either a quantity.Energy or a quantity.Force specifying whether
energy_or_force is an energy or force respectively.
dt_start: The initial step size during minimization as a float.
dt_max: The maximum step size during minimization as a float.
n_min: An integer specifying the minimum number of steps moving in the
correct direction before dt and f_alpha should be updated.
f_inc: A float specifying the fractional rate by which the step size
should be increased.
f_dec: A float specifying the fractional rate by which the step size
should be decreased.
alpha_start: A float specifying the initial momentum.
f_alpha: A float specifying the fractional change in momentum.
Returns:
See above.
[1] Bitzek, Erik, Pekka Koskinen, Franz Gahler, Michael Moseler,
and Peter Gumbsch. "Structural relaxation made simple."
Physical review letters 97, no. 17 (2006): 170201.
"""
dt_start, dt_max, n_min, f_inc, f_dec, alpha_start, f_alpha = static_cast(
dt_start, dt_max, n_min, f_inc, f_dec, alpha_start, f_alpha)
force = quantity.canonicalize_force(energy_or_force, quant)
def init_fun(R, **kwargs):
V = np.zeros_like(R)
return FireDescentState(R, V, force(R, **kwargs), dt_start,
alpha_start, f32(0))
def apply_fun(state, **kwargs):
R, V, F_old, dt, alpha, n_pos = state
R = shift_fn(R, dt * V + dt**f32(2) * F_old, **kwargs)
F = force(R, **kwargs)
V = V + dt * f32(0.5) * (F_old + F)
F_norm = np.sqrt(np.sum(F**f32(2)) + f32(1e-6))
V_norm = np.sqrt(np.sum(V**f32(2)))
P = np.array(np.dot(np.reshape(F, (-1)), np.reshape(V, (-1))))
V = V + alpha * (F * V_norm / F_norm - V)
n_pos = np.where(P >= 0, n_pos + f32(1.0), f32(0))
dt_choice = np.array([dt * f_inc, dt_max])
dt = np.where(P > 0, np.where(n_pos > n_min, np.min(dt_choice), dt),
dt)
dt = np.where(P < 0, dt * f_dec, dt)
alpha = np.where(P > 0, np.where(n_pos > n_min, alpha * f_alpha,
alpha), alpha)
alpha = np.where(P < 0, alpha_start, alpha)
V = (P < 0) * np.zeros_like(V) + (P >= 0) * V
return FireDescentState(R, V, F, dt, alpha, n_pos)
return init_fun, apply_fun
def nvt_nose_hoover(energy_or_force_fn: Callable[..., Array],
shift_fn: ShiftFn,
dt: float,
kT: float,
chain_length: int = 5,
chain_steps: int = 2,
sy_steps: int = 3,
tau: Optional[float] = None) -> Simulator:
"""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 (NHC) thermostat described in [1, 2, 3]. We follow the
direct translation method outlined in [3] and the interested reader might
want to look at that paper as a reference.
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.
kT: Floating point number specifying the temperature inunits of Boltzmann
constant. To update the temperature dynamically during a simulation one
should pass `kT` as a keyword argument to the step function.
chain_length: An integer specifying the number of particles in
the Nose-Hoover chain.
chain_steps: An integer specifying the number, $n_c$, of outer substeps.
sy_steps: An integer specifying the number of Suzuki-Yoshida steps. This
must be either 1, 3, 5, or 7.
tau: A floating point timescale over which temperature equilibration occurs.
Measured in units of dt. The performance of the Nose-Hoover chain
thermostat can be 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.
"""
dt = f32(dt)
if tau is None:
tau = dt * 100
tau = f32(tau)
force_fn = quantity.canonicalize_force(energy_or_force_fn)
chain_fns = nose_hoover_chain(dt, chain_length, chain_steps, sy_steps, tau)
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 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
return init_fn, apply_fn
