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