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