Last active
August 29, 2015 14:01
-
-
Save synapticarbors/fb188121b0a3891109f1 to your computer and use it in GitHub Desktop.
Modified WCA-dimer script that uses CustomIntegrator for GHMC propagation. The Integrator module is available at: https://github.com/choderalab/openmm-tests/blob/master/energy-rms/integrators.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #============================================================================================= | |
| # MODULE DOCSTRING | |
| #============================================================================================= | |
| """ | |
| Custom integrators. | |
| DESCRIPTION | |
| TODO | |
| COPYRIGHT AND LICENSE | |
| @author John D. Chodera <jchodera@gmail.com> | |
| All code in this repository is released under the GNU General Public License. | |
| This program is free software: you can redistribute it and/or modify it under | |
| the terms of the GNU General Public License as published by the Free Software | |
| Foundation, either version 3 of the License, or (at your option) any later | |
| version. | |
| This program is distributed in the hope that it will be useful, but WITHOUT ANY | |
| WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A | |
| PARTICULAR PURPOSE. See the GNU General Public License for more details. | |
| You should have received a copy of the GNU General Public License along with | |
| this program. If not, see <http://www.gnu.org/licenses/>. | |
| """ | |
| #============================================================================================= | |
| # GLOBAL IMPORTS | |
| #============================================================================================= | |
| import numpy | |
| import simtk.unit | |
| import simtk.unit as units | |
| import simtk.openmm as mm | |
| #============================================================================================= | |
| # CONSTANTS | |
| #============================================================================================= | |
| kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA | |
| #============================================================================================= | |
| # INTEGRATORS | |
| #============================================================================================= | |
| def DummyIntegrator(): | |
| """ | |
| Construct a dummy integrator that does nothing except update call the force updates. | |
| Returns | |
| ------- | |
| integrator : simtk.openmm.CustomIntegrator | |
| A dummy integrator. | |
| Examples | |
| -------- | |
| Create a dummy integrator. | |
| >>> integrator = DummyIntegrator() | |
| """ | |
| timestep = 0.0 * units.femtoseconds | |
| integrator = mm.CustomIntegrator(timestep) | |
| integrator.addUpdateContextState() | |
| integrator.addConstrainPositions() | |
| integrator.addConstrainVelocities() | |
| return integrator | |
| def GradientDescentMinimizationIntegrator(initial_step_size=0.01*units.angstroms): | |
| """ | |
| Construct a simple gradient descent minimization integrator. | |
| Parameters | |
| ---------- | |
| initial_step_size : numpy.unit.Quantity compatible with nanometers, default: 0.01*simtk.unit.angstroms | |
| The norm of the initial step size guess. | |
| Returns | |
| ------- | |
| integrator : simtk.openmm.CustomIntegrator | |
| A velocity Verlet integrator. | |
| Notes | |
| ----- | |
| An adaptive step size is used. | |
| References | |
| ---------- | |
| Examples | |
| -------- | |
| Create a gradient descent minimization integrator. | |
| >>> integrator = GradientDescentMinimizationIntegrator() | |
| """ | |
| timestep = 1.0 * units.femtoseconds | |
| integrator = mm.CustomIntegrator(timestep) | |
| integrator.addGlobalVariable("step_size", initial_step_size/units.nanometers) | |
| integrator.addGlobalVariable("energy_old", 0) | |
| integrator.addGlobalVariable("energy_new", 0) | |
| integrator.addGlobalVariable("delta_energy", 0) | |
| integrator.addGlobalVariable("accept", 0) | |
| integrator.addGlobalVariable("fnorm2", 0) | |
| integrator.addPerDofVariable("x_old", 0) | |
| # Update context state. | |
| integrator.addUpdateContextState() | |
| # Constrain positions. | |
| integrator.addConstrainPositions() | |
| # Store old energy and positions. | |
| integrator.addComputeGlobal("energy_old", "energy") | |
| integrator.addComputePerDof("x_old", "x") | |
| # Compute sum of squared norm. | |
| integrator.addComputeSum("fnorm2", "f^2") | |
| # Take step. | |
| integrator.addComputePerDof("x", "x+step_size*f/sqrt(fnorm2 + delta(fnorm2))") | |
| integrator.addConstrainPositions() | |
| # Ensure we only keep steps that go downhill in energy. | |
| integrator.addComputeGlobal("energy_new", "energy") | |
| integrator.addComputeGlobal("delta_energy", "energy_new-energy_old") | |
| # Accept also checks for NaN | |
| integrator.addComputeGlobal("accept", "step(-delta_energy) * delta(energy - energy_new)") | |
| integrator.addComputePerDof("x", "accept*x + (1-accept)*x_old") | |
| # Update step size. | |
| integrator.addComputeGlobal("step_size", "step_size * (2.0*accept + 0.5*(1-accept))") | |
| return integrator | |
| def VelocityVerletIntegrator(timestep=1.0*simtk.unit.femtoseconds): | |
| """ | |
| Construct a velocity Verlet integrator. | |
| Parameters | |
| ---------- | |
| timestep : numpy.unit.Quantity compatible with femtoseconds, default: 1*simtk.unit.femtoseconds | |
| The integration timestep. | |
| Returns | |
| ------- | |
| integrator : simtk.openmm.CustomIntegrator | |
| A velocity Verlet integrator. | |
| Notes | |
| ----- | |
| This integrator is taken verbatim from Peter Eastman's example appearing in the CustomIntegrator header file documentation. | |
| References | |
| ---------- | |
| W. C. Swope, H. C. Andersen, P. H. Berens, and K. R. Wilson, J. Chem. Phys. 76, 637 (1982) | |
| Examples | |
| -------- | |
| Create a velocity Verlet integrator. | |
| >>> timestep = 1.0 * simtk.unit.femtoseconds | |
| >>> integrator = VelocityVerletIntegrator(timestep) | |
| """ | |
| integrator = mm.CustomIntegrator(timestep) | |
| integrator.addPerDofVariable("x1", 0) | |
| integrator.addUpdateContextState() | |
| integrator.addComputePerDof("v", "v+0.5*dt*f/m") | |
| integrator.addComputePerDof("x", "x+dt*v") | |
| integrator.addComputePerDof("x1", "x") | |
| integrator.addConstrainPositions() | |
| integrator.addComputePerDof("v", "v+0.5*dt*f/m+(x-x1)/dt") | |
| integrator.addConstrainVelocities() | |
| return integrator | |
| def AndersenVelocityVerletIntegrator(temperature=298*simtk.unit.kelvin, collision_rate=91.0/simtk.unit.picoseconds, timestep=1.0*simtk.unit.femtoseconds): | |
| """ | |
| Construct a velocity Verlet integrator with Andersen thermostat, implemented as per-particle collisions (rather than massive collisions). | |
| Parameters | |
| ---------- | |
| temperature : numpy.unit.Quantity compatible with kelvin, default: 298*simtk.unit.kelvin | |
| The temperature of the fictitious bath. | |
| collision_rate : numpy.unit.Quantity compatible with 1/picoseconds, default: 91/simtk.unit.picoseconds | |
| The collision rate with fictitious bath particles. | |
| timestep : numpy.unit.Quantity compatible with femtoseconds, default: 1*simtk.unit.femtoseconds | |
| The integration timestep. | |
| Returns | |
| ------- | |
| integrator : simtk.openmm.CustomIntegrator | |
| A velocity Verlet integrator with periodic Andersen thermostat. | |
| References | |
| ---------- | |
| Hans C. Andersen "Molecular dynamics simulations at constant pressure and/or temperature", Journal of Chemical Physics 72, 2384-2393 (1980) | |
| http://dx.doi.org/10.1063/1.439486 | |
| Examples | |
| -------- | |
| Create a velocity Verlet integrator with Andersen thermostat. | |
| >>> timestep = 1.0 * simtk.unit.femtoseconds | |
| >>> collision_rate = 91.0 / simtk.unit.picoseconds | |
| >>> temperature = 298.0 * simtk.unit.kelvin | |
| >>> integrator = AndersenVelocityVerletIntegrator(temperature, collision_rate, timestep) | |
| Notes | |
| ------ | |
| The velocity Verlet integrator is taken verbatim from Peter Eastman's example in the CustomIntegrator header file documentation. | |
| The efficiency could be improved by avoiding recomputation of sigma_v every timestep. | |
| """ | |
| integrator = mm.CustomIntegrator(timestep) | |
| # | |
| # Integrator initialization. | |
| # | |
| kT = kB * temperature | |
| integrator.addGlobalVariable("kT", kT) # thermal energy | |
| integrator.addGlobalVariable("p_collision", timestep * collision_rate) # per-particle collision probability per timestep | |
| integrator.addPerDofVariable("sigma_v", 0) # velocity distribution stddev for Maxwell-Boltzmann (computed later) | |
| integrator.addPerDofVariable("collision", 0) # 1 if collision has occured this timestep, 0 otherwise | |
| integrator.addPerDofVariable("x1", 0) # for constraints | |
| # | |
| # Update velocities from Maxwell-Boltzmann distribution for particles that collide. | |
| # | |
| integrator.addComputePerDof("sigma_v", "sqrt(kT/m)") | |
| integrator.addComputePerDof("collision", "step(p_collision-uniform)") # if collision has occured this timestep, 0 otherwise | |
| integrator.addComputePerDof("v", "(1-collision)*v + collision*sigma_v*gaussian") # randomize velocities of particles that have collided | |
| # | |
| # Velocity Verlet step | |
| # | |
| integrator.addUpdateContextState() | |
| integrator.addComputePerDof("v", "v+0.5*dt*f/m") | |
| integrator.addComputePerDof("x", "x+dt*v") | |
| integrator.addComputePerDof("x1", "x") | |
| integrator.addConstrainPositions() | |
| integrator.addComputePerDof("v", "v+0.5*dt*f/m+(x-x1)/dt") | |
| integrator.addConstrainVelocities() | |
| return integrator | |
| def MetropolisMonteCarloIntegrator(temperature=298.0*simtk.unit.kelvin, sigma=0.1*simtk.unit.angstroms, timestep=1*simtk.unit.femtoseconds): | |
| """ | |
| Create a simple Metropolis Monte Carlo integrator that uses Gaussian displacement trials. | |
| Parameters | |
| ---------- | |
| temperature : numpy.unit.Quantity compatible with kelvin, default: 298*simtk.unit.kelvin | |
| The temperature. | |
| sigma : numpy.unit.Quantity compatible with nanometers, default: 0.1*simtk.unit.angstroms | |
| The displacement standard deviation for each degree of freedom. | |
| timestep : numpy.unit.Quantity compatible with femtoseconds, default: 1.0*simtk.unit.femtoseconds | |
| The integration timestep, which is purely fictitious---it is just used to advance the simulation clock. | |
| Returns | |
| ------- | |
| integrator : simtk.openmm.CustomIntegrator | |
| A Metropolis Monte Carlo integrator. | |
| Warning | |
| ------- | |
| This integrator does not respect constraints. | |
| Notes | |
| ----- | |
| The timestep is purely fictitious, and just used to advance the simulation clock. | |
| Velocities are drawn from a Maxwell-Boltzmann distribution each timestep to generate correct (x,v) statistics. | |
| Additional global variables 'ntrials' and 'naccept' keep track of how many trials have been attempted and accepted, respectively. | |
| Examples | |
| -------- | |
| Create a Metropolis Monte Carlo integrator with specified random displacement standard deviation. | |
| >>> timestep = 1.0 * simtk.unit.femtoseconds # fictitious timestep | |
| >>> temperature = 298.0 * simtk.unit.kelvin | |
| >>> sigma = 1.0 * simtk.unit.angstroms | |
| >>> integrator = MetropolisMonteCarloIntegrator(temperature, sigma, timestep) | |
| """ | |
| # Create a new Custom integrator. | |
| integrator = mm.CustomIntegrator(timestep) | |
| # Compute the thermal energy. | |
| kT = kB * temperature | |
| # | |
| # Integrator initialization. | |
| # | |
| integrator.addGlobalVariable("naccept", 0) # number accepted | |
| integrator.addGlobalVariable("ntrials", 0) # number of Metropolization trials | |
| integrator.addGlobalVariable("kT", kT) # thermal energy | |
| integrator.addPerDofVariable("sigma_x", sigma) # perturbation size | |
| integrator.addPerDofVariable("sigma_v", 0) # velocity distribution stddev for Maxwell-Boltzmann (set later) | |
| integrator.addPerDofVariable("xold", 0) # old positions | |
| integrator.addGlobalVariable("Eold", 0) # old energy | |
| integrator.addGlobalVariable("Enew", 0) # new energy | |
| integrator.addGlobalVariable("accept", 0) # accept or reject | |
| # | |
| # Context state update. | |
| # | |
| integrator.addUpdateContextState(); | |
| # | |
| # Update velocities from Maxwell-Boltzmann distribution. | |
| # | |
| integrator.addComputePerDof("sigma_v", "sqrt(kT/m)") | |
| integrator.addComputePerDof("v", "sigma_v*gaussian") | |
| integrator.addConstrainVelocities(); | |
| # | |
| # propagation steps | |
| # | |
| # Store old positions and energy. | |
| integrator.addComputePerDof("xold", "x") | |
| integrator.addComputeGlobal("Eold", "energy") | |
| # Gaussian particle displacements. | |
| integrator.addComputePerDof("x", "x + sigma_x*gaussian") | |
| # Accept or reject with Metropolis criteria. | |
| integrator.addComputeGlobal("accept", "step(exp(-(energy-Eold)/kT) - uniform)") | |
| integrator.addComputePerDof("x", "(1-accept)*xold + x*accept") | |
| # Accumulate acceptance statistics. | |
| integrator.addComputeGlobal("naccept", "naccept + accept") | |
| integrator.addComputeGlobal("ntrials", "ntrials + 1") | |
| return integrator | |
| def HMCIntegrator(temperature=298.0*simtk.unit.kelvin, nsteps=10, timestep=1*simtk.unit.femtoseconds): | |
| """ | |
| Create a hybrid Monte Carlo (HMC) integrator. | |
| Parameters | |
| ---------- | |
| temperature : numpy.unit.Quantity compatible with kelvin, default: 298*simtk.unit.kelvin | |
| The temperature. | |
| nsteps : int, default: 10 | |
| The number of velocity Verlet steps to take per HMC trial. | |
| timestep : numpy.unit.Quantity compatible with femtoseconds, default: 1*simtk.unit.femtoseconds | |
| The integration timestep. | |
| Returns | |
| ------- | |
| integrator : simtk.openmm.CustomIntegrator | |
| A hybrid Monte Carlo integrator. | |
| Warning | |
| ------- | |
| Because 'nsteps' sets the number of steps taken, a call to integrator.step(1) actually takes 'nsteps' steps. | |
| Notes | |
| ----- | |
| The velocity is drawn from a Maxwell-Boltzmann distribution, then 'nsteps' steps are taken, | |
| and the new configuration is either accepted or rejected. | |
| Additional global variables 'ntrials' and 'naccept' keep track of how many trials have been attempted and | |
| accepted, respectively. | |
| TODO | |
| ---- | |
| Currently, the simulation timestep is only advanced by 'timestep' each step, rather than timestep*nsteps. Fix this. | |
| Examples | |
| -------- | |
| Create an HMC integrator. | |
| >>> timestep = 1.0 * simtk.unit.femtoseconds # fictitious timestep | |
| >>> temperature = 298.0 * simtk.unit.kelvin | |
| >>> nsteps = 10 # number of steps per call | |
| >>> integrator = HMCIntegrator(temperature, nsteps, timestep) | |
| """ | |
| # Create a new custom integrator. | |
| integrator = mm.CustomIntegrator(timestep) | |
| # Compute the thermal energy. | |
| kT = kB * temperature | |
| # | |
| # Integrator initialization. | |
| # | |
| integrator.addGlobalVariable("naccept", 0) # number accepted | |
| integrator.addGlobalVariable("ntrials", 0) # number of Metropolization trials | |
| integrator.addGlobalVariable("kT", kT) # thermal energy | |
| integrator.addPerDofVariable("sigma", 0) | |
| integrator.addGlobalVariable("ke", 0) # kinetic energy | |
| integrator.addPerDofVariable("xold", 0) # old positions | |
| integrator.addGlobalVariable("Eold", 0) # old energy | |
| integrator.addGlobalVariable("Enew", 0) # new energy | |
| integrator.addGlobalVariable("accept", 0) # accept or reject | |
| integrator.addPerDofVariable("x1", 0) # for constraints | |
| # | |
| # Pre-computation. | |
| # This only needs to be done once, but it needs to be done for each degree of freedom. | |
| # Could move this to initialization? | |
| # | |
| integrator.addComputePerDof("sigma", "sqrt(kT/m)") | |
| # | |
| # Allow Context updating here, outside of inner loop only. | |
| # | |
| integrator.addUpdateContextState(); | |
| # | |
| # Draw new velocity. | |
| # | |
| integrator.addComputePerDof("v", "sigma*gaussian") | |
| integrator.addConstrainVelocities(); | |
| # | |
| # Store old position and energy. | |
| # | |
| integrator.addComputeSum("ke", "0.5*m*v*v") | |
| integrator.addComputeGlobal("Eold", "ke + energy") | |
| integrator.addComputePerDof("xold", "x") | |
| # | |
| # Inner symplectic steps using velocity Verlet. | |
| # | |
| for step in range(nsteps): | |
| integrator.addComputePerDof("v", "v+0.5*dt*f/m") | |
| integrator.addComputePerDof("x", "x+dt*v") | |
| integrator.addComputePerDof("x1", "x") | |
| integrator.addConstrainPositions() | |
| integrator.addComputePerDof("v", "v+0.5*dt*f/m+(x-x1)/dt") | |
| integrator.addConstrainVelocities() | |
| # | |
| # Accept/reject step. | |
| # | |
| integrator.addComputeSum("ke", "0.5*m*v*v") | |
| integrator.addComputeGlobal("Enew", "ke + energy") | |
| integrator.addComputeGlobal("accept", "step(exp(-(Enew-Eold)/kT) - uniform)") | |
| integrator.addComputePerDof("x", "x*accept + xold*(1-accept)") | |
| # | |
| # Accumulate statistics. | |
| # | |
| integrator.addComputeGlobal("naccept", "naccept + accept") | |
| integrator.addComputeGlobal("ntrials", "ntrials + 1") | |
| return integrator | |
| def GHMCIntegrator(temperature=298.0*simtk.unit.kelvin, collision_rate=91.0/simtk.unit.picoseconds, timestep=1.0*simtk.unit.femtoseconds): | |
| """ | |
| Create a generalized hybrid Monte Carlo (GHMC) integrator. | |
| Parameters | |
| ---------- | |
| temperature : numpy.unit.Quantity compatible with kelvin, default: 298*simtk.unit.kelvin | |
| The temperature. | |
| collision_rate : numpy.unit.Quantity compatible with 1/picoseconds, default: 91.0/simtk.unit.picoseconds | |
| The collision rate. | |
| timestep : numpy.unit.Quantity compatible with femtoseconds, default: 1.0*simtk.unit.femtoseconds | |
| The integration timestep. | |
| Returns | |
| ------- | |
| integrator : simtk.openmm.CustomIntegrator | |
| A GHMC integrator. | |
| Notes | |
| ----- | |
| This integrator is equivalent to a Langevin integrator in the velocity Verlet discretization with a | |
| Metrpolization step to ensure sampling from the appropriate distribution. | |
| Additional global variables 'ntrials' and 'naccept' keep track of how many trials have been attempted and | |
| accepted, respectively. | |
| TODO | |
| ---- | |
| Move initialization of 'sigma' to setting the per-particle variables. | |
| Examples | |
| -------- | |
| Create a GHMC integrator. | |
| >>> temperature = 298.0 * simtk.unit.kelvin | |
| >>> collision_rate = 91.0 / simtk.unit.picoseconds | |
| >>> timestep = 1.0 * simtk.unit.femtoseconds | |
| >>> integrator = GHMCIntegrator(temperature, collision_rate, timestep) | |
| References | |
| ---------- | |
| Lelievre T, Stoltz G, and Rousset M. Free Energy Computations: A Mathematical Perspective | |
| http://www.amazon.com/Free-Energy-Computations-Mathematical-Perspective/dp/1848162472 | |
| """ | |
| # Initialize constants. | |
| kT = kB * temperature | |
| gamma = collision_rate | |
| # Create a new custom integrator. | |
| integrator = mm.CustomIntegrator(timestep) | |
| # | |
| # Integrator initialization. | |
| # | |
| integrator.addGlobalVariable("kT", kT) # thermal energy | |
| integrator.addGlobalVariable("b", numpy.exp(-gamma*timestep)) # velocity mixing parameter | |
| integrator.addPerDofVariable("sigma", 0) | |
| integrator.addGlobalVariable("ke", 0) # kinetic energy | |
| integrator.addPerDofVariable("vold", 0) # old velocities | |
| integrator.addPerDofVariable("xold", 0) # old positions | |
| integrator.addGlobalVariable("Eold", 0) # old energy | |
| integrator.addGlobalVariable("Enew", 0) # new energy | |
| integrator.addGlobalVariable("accept", 0) # accept or reject | |
| integrator.addGlobalVariable("naccept", 0) # number accepted | |
| integrator.addGlobalVariable("ntrials", 0) # number of Metropolization trials | |
| integrator.addPerDofVariable("x1", 0) # position before application of constraints | |
| # | |
| # Pre-computation. | |
| # This only needs to be done once, but it needs to be done for each degree of freedom. | |
| # Could move this to initialization? | |
| # | |
| integrator.addComputePerDof("sigma", "sqrt(kT/m)") | |
| # | |
| # Allow context updating here. | |
| # | |
| integrator.addUpdateContextState(); | |
| # | |
| # Constrain positions. | |
| # | |
| integrator.addConstrainPositions(); | |
| # | |
| # Velocity perturbation. | |
| # | |
| integrator.addComputePerDof("v", "sqrt(b)*v + sqrt(1-b)*sigma*gaussian") | |
| integrator.addConstrainVelocities(); | |
| # | |
| # Metropolized symplectic step. | |
| # | |
| integrator.addComputeSum("ke", "0.5*m*v*v") | |
| integrator.addComputeGlobal("Eold", "ke + energy") | |
| integrator.addComputePerDof("xold", "x") | |
| integrator.addComputePerDof("vold", "v") | |
| integrator.addComputePerDof("v", "v + 0.5*dt*f/m") | |
| integrator.addComputePerDof("x", "x + v*dt") | |
| integrator.addComputePerDof("x1", "x") | |
| integrator.addConstrainPositions(); | |
| integrator.addComputePerDof("v", "v + 0.5*dt*f/m + (x-x1)/dt") | |
| integrator.addConstrainVelocities(); | |
| integrator.addComputeSum("ke", "0.5*m*v*v") | |
| integrator.addComputeGlobal("Enew", "ke + energy") | |
| integrator.addComputeGlobal("accept", "step(exp(-(Enew-Eold)/kT) - uniform)") | |
| integrator.addComputePerDof("x", "x*accept + xold*(1-accept)") | |
| integrator.addComputePerDof("v", "v*accept - vold*(1-accept)") | |
| # | |
| # Velocity randomization | |
| # | |
| integrator.addComputePerDof("v", "sqrt(b)*v + sqrt(1-b)*sigma*gaussian") | |
| integrator.addConstrainVelocities(); | |
| # | |
| # Accumulate statistics. | |
| # | |
| integrator.addComputeGlobal("naccept", "naccept + accept") | |
| integrator.addComputeGlobal("ntrials", "ntrials + 1") | |
| return integrator | |
| def VVVRIntegrator(temperature=298.0*simtk.unit.kelvin, collision_rate=91.0/simtk.unit.picoseconds, timestep=1.0*simtk.unit.femtoseconds): | |
| """ | |
| Create a velocity verlet with velocity randomization (VVVR) integrator. | |
| Parameters | |
| ---------- | |
| temperature : numpy.unit.Quantity compatible with kelvin, default: 298.0*simtk.unit.kelvin | |
| The temperature. | |
| collision_rate : numpy.unit.Quantity compatible with 1/picoseconds, default: 91.0/simtk.unit.picoseconds | |
| The collision rate. | |
| timestep : numpy.unit.Quantity compatible with femtoseconds, default: 1.0*simtk.unit.femtoseconds | |
| The integration timestep. | |
| Returns | |
| ------- | |
| integrator : simtk.openmm.CustomIntegrator | |
| VVVR integrator. | |
| Notes | |
| ----- | |
| This integrator is equivalent to a Langevin integrator in the velocity Verlet discretization with a | |
| timestep correction to ensure that the field-free diffusion constant is timestep invariant. | |
| The global 'pseudowork' keeps track of the pseudowork accumulated during integration, and can be | |
| used to correct the sampled statistics or in a Metropolization scheme. | |
| TODO | |
| ---- | |
| Move initialization of 'sigma' to setting the per-particle variables. | |
| We can ditch pseudowork and instead use total energy difference - heat. | |
| References | |
| ---------- | |
| David A. Sivak, John D. Chodera, and Gavin E. Crooks. | |
| Time step rescaling recovers continuous-time dynamical properties for discrete-time Langevin integration of nonequilibrium systems | |
| http://arxiv.org/abs/1301.3800 | |
| Examples | |
| -------- | |
| Create a VVVR integrator. | |
| >>> temperature = 298.0 * simtk.unit.kelvin | |
| >>> collision_rate = 91.0 / simtk.unit.picoseconds | |
| >>> timestep = 1.0 * simtk.unit.femtoseconds | |
| >>> integrator = VVVRIntegrator(temperature, collision_rate, timestep) | |
| """ | |
| # Compute constants. | |
| kT = kB * temperature | |
| gamma = collision_rate | |
| # Create a new custom integrator. | |
| integrator = mm.CustomIntegrator(timestep) | |
| # | |
| # Integrator initialization. | |
| # | |
| integrator.addGlobalVariable("kT", kT) # thermal energy | |
| integrator.addGlobalVariable("b", numpy.exp(-gamma*timestep)) # velocity mixing parameter | |
| integrator.addPerDofVariable("sigma", 0) | |
| integrator.addPerDofVariable("x1", 0) # position before application of constraints | |
| # | |
| # Allow context updating here. | |
| # | |
| integrator.addUpdateContextState(); | |
| # | |
| # Pre-computation. | |
| # This only needs to be done once, but it needs to be done for each degree of freedom. | |
| # Could move this to initialization? | |
| # | |
| integrator.addComputePerDof("sigma", "sqrt(kT/m)") | |
| # | |
| # Velocity perturbation. | |
| # | |
| integrator.addComputePerDof("v", "sqrt(b)*v + sqrt(1-b)*sigma*gaussian") | |
| integrator.addConstrainVelocities(); | |
| # | |
| # Metropolized symplectic step. | |
| # | |
| integrator.addComputePerDof("v", "v + 0.5*dt*f/m") | |
| integrator.addComputePerDof("x", "x + v*dt") | |
| integrator.addComputePerDof("x1", "x") | |
| integrator.addConstrainPositions(); | |
| integrator.addComputePerDof("v", "v + 0.5*dt*f/m + (x-x1)/dt") | |
| integrator.addConstrainVelocities(); | |
| # | |
| # Velocity randomization | |
| # | |
| integrator.addComputePerDof("v", "sqrt(b)*v + sqrt(1-b)*sigma*gaussian") | |
| integrator.addConstrainVelocities(); | |
| return integrator | |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #============================================================================================= | |
| # MODULE DOCSTRING | |
| #============================================================================================= | |
| """ | |
| Simulation of WCA dimer in dense WCA solvent using GHMC. | |
| DESCRIPTION | |
| COPYRIGHT | |
| @author John D. Chodera <jchodera@gmail.com> | |
| All code in this repository is released under the GNU General Public License. | |
| This program is free software: you can redistribute it and/or modify it under | |
| the terms of the GNU General Public License as published by the Free Software | |
| Foundation, either version 3 of the License, or (at your option) any later | |
| version. | |
| This program is distributed in the hope that it will be useful, but WITHOUT ANY | |
| WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A | |
| PARTICULAR PURPOSE. See the GNU General Public License for more details. | |
| You should have received a copy of the GNU General Public License along with | |
| this program. If not, see <http://www.gnu.org/licenses/>. | |
| TODO | |
| """ | |
| #============================================================================================= | |
| # GLOBAL IMPORTS | |
| #============================================================================================= | |
| import os | |
| import os.path | |
| import sys | |
| import math | |
| import copy | |
| import time | |
| import numpy | |
| import simtk | |
| import simtk.unit as units | |
| import simtk.openmm as openmm | |
| #import Scientific.IO.NetCDF as netcdf # for netcdf interface in Scientific | |
| import netCDF4 as netcdf # for netcdf interface provided by netCDF4 in enthought python | |
| import wcadimer | |
| import sampling | |
| from integrators import GHMCIntegrator | |
| #============================================================================================= | |
| # SUBROUTINES | |
| #============================================================================================= | |
| def norm(n01): | |
| return n01.unit * numpy.sqrt(numpy.dot(n01/n01.unit, n01/n01.unit)) | |
| #============================================================================================= | |
| # MAIN AND TESTS | |
| #============================================================================================= | |
| if __name__ == "__main__": | |
| # PARAMETERS | |
| netcdf_filename = 'data/md-solvent.nc' | |
| verbose = False | |
| # WCA fluid parameters (argon). | |
| mass = wcadimer.mass | |
| sigma = wcadimer.sigma | |
| epsilon = wcadimer.epsilon | |
| r_WCA = wcadimer.r_WCA | |
| r0 = wcadimer.r0 | |
| h = wcadimer.h | |
| w = wcadimer.w | |
| # Compute characteristic timescale. | |
| tau = wcadimer.tau | |
| print "tau = %.3f ps" % (tau / units.picoseconds) | |
| # Compute timestep. | |
| equilibrate_timestep = 2 * wcadimer.stable_timestep | |
| timestep = 5 * wcadimer.stable_timestep | |
| print "equilibrate timestep = %.1f fs, switch timestep = %.1f fs" % (equilibrate_timestep / units.femtoseconds, timestep / units.femtoseconds) | |
| # Set temperature, pressure, and collision rate for stochastic thermostats. | |
| temperature = wcadimer.temperature | |
| print "temperature = %.1f K" % (temperature / units.kelvin) | |
| kT = wcadimer.kT | |
| beta = 1.0 / kT # inverse temperature | |
| collision_rate = 1.0 / tau # collision rate for Langevin integrator | |
| print 'collision_rate: ', collision_rate | |
| niterations = 10 # number of work samples to collect | |
| # Create system. | |
| [system, coordinates] = wcadimer.WCADimer() | |
| # Form vectors of masses and sqrt(kT/m) for force propagation and velocity randomization. | |
| print "Creating masses array..." | |
| nparticles = system.getNumParticles() | |
| masses = units.Quantity(numpy.zeros([nparticles,3], numpy.float64), units.amu) | |
| for particle_index in range(nparticles): | |
| masses[particle_index,:] = system.getParticleMass(particle_index) | |
| sqrt_kT_over_m = units.Quantity(numpy.zeros([nparticles,3], numpy.float64), units.nanometers / units.picosecond) | |
| for particle_index in range(nparticles): | |
| sqrt_kT_over_m[particle_index,:] = units.sqrt(kT / masses[particle_index,0]) # standard deviation of velocity distribution for each coordinate for this atom | |
| # List all available platforms | |
| print "Available platforms:" | |
| for platform_index in range(openmm.Platform.getNumPlatforms()): | |
| platform = openmm.Platform.getPlatform(platform_index) | |
| print "%5d %s" % (platform_index, platform.getName()) | |
| print "" | |
| # Select platform. | |
| #platform = openmm.Platform.getPlatformByName("CPU") | |
| platform = openmm.Platform.getPlatformByName("CPU") | |
| min_platform = openmm.Platform.getPlatformByName("Reference") | |
| #deviceid = 3 | |
| #platform.setPropertyDefaultValue('OpenCLDeviceIndex', '%d' % deviceid) | |
| #platform.setPropertyDefaultValue('CudaDeviceIndex', '%d' % deviceid) | |
| #platform.setPropertyDefaultValue('OpenCLPrecision', 'mixed') | |
| #platform.setPropertyDefaultValue('CudaPrecision', 'double') | |
| # Initialize netcdf file. | |
| if not os.path.exists(netcdf_filename): | |
| # Open NetCDF file for writing | |
| ncfile = netcdf.Dataset(netcdf_filename, 'w') # for netCDF4 | |
| # Create dimensions. | |
| ncfile.createDimension('iteration', 0) # unlimited number of iterations | |
| ncfile.createDimension('nparticles', nparticles) # number of particles | |
| ncfile.createDimension('ndim', 3) # number of dimensions | |
| # Create variables. | |
| ncfile.createVariable('distance', 'd', ('iteration',)) | |
| ncfile.createVariable('positions', 'd', ('iteration','nparticles','ndim')) | |
| ncfile.createVariable('fraction_accepted', 'd', ('iteration',)) | |
| # Force sync to disk to avoid data loss. | |
| ncfile.sync() | |
| # Minimize. | |
| print "Minimizing energy..." | |
| coordinates = sampling.minimize(min_platform, system, coordinates) | |
| # Equilibrate. | |
| print "Equilibrating..." | |
| [coordinates, velocities, fraction_accepted] = sampling.equilibrate_ghmc(system, equilibrate_timestep, collision_rate, temperature, masses, sqrt_kT_over_m, coordinates, platform) | |
| print "%.3f %% accepted" % (fraction_accepted * 100.0) | |
| # Write initial configuration. | |
| ncfile.variables['distance'][0] = norm(coordinates[1,:] - coordinates[0,:]) / units.angstroms | |
| ncfile.variables['positions'][0,:,:] = coordinates[:,:] / units.angstroms | |
| ncfile.sync() | |
| iteration = 1 | |
| else: | |
| # Open NetCDF file for reading. | |
| ncfile = netcdf.Dataset(netcdf_filename, 'a') # for netCDF4 | |
| # Read iteration and coordinates. | |
| iteration = ncfile.variables['distance'][:].size | |
| coordinates = units.Quantity(ncfile.variables['positions'][iteration-1,:,:], units.angstroms) | |
| # Continue | |
| print timestep.in_units_of(units.femtosecond) | |
| print temperature.in_units_of(units.kelvin) | |
| print collision_rate.in_units_of(units.picosecond**-1) | |
| ghmc_integrator = GHMCIntegrator(timestep=timestep, temperature=temperature, | |
| collision_rate=collision_rate) | |
| ghmc_global_variables = { ghmc_integrator.getGlobalVariableName(index) : index for index in range(ghmc_integrator.getNumGlobalVariables()) } | |
| context = openmm.Context(system, ghmc_integrator, platform) | |
| context.setPositions(coordinates) | |
| state = context.getState(getPositions=True) | |
| while (iteration < niterations): | |
| print "iteration %5d / %5d" % (iteration, niterations) | |
| print 'coordinates: ', coordinates[0,:] / units.angstrom | |
| initial_time = time.time() | |
| # Generate a new configuration. | |
| initial_distance = norm(coordinates[1,:] - coordinates[0,:]) | |
| ghmc_integrator.setGlobalVariable(ghmc_global_variables['naccept'], 0) | |
| ghmc_integrator.setGlobalVariable(ghmc_global_variables['ntrials'], 0) | |
| ghmc_integrator.step(500) | |
| naccept = ghmc_integrator.getGlobalVariable(ghmc_global_variables['naccept']) | |
| ntrials = ghmc_integrator.getGlobalVariable(ghmc_global_variables['ntrials']) | |
| fraction_accepted = float(naccept) / float(ntrials) | |
| print "%.3f %% accepted" % (fraction_accepted * 100.0) | |
| state = context.getState(getPositions=True) | |
| coordinates = state.getPositions(asNumpy=True) | |
| final_distance = norm(coordinates[1,:] - coordinates[0,:]) | |
| print "Dynamics %.1f A -> %.1f A (barrier at %.1f A)" % (initial_distance / units.angstroms, final_distance / units.angstroms, (r0+w)/units.angstroms) | |
| ncfile.variables['fraction_accepted'][iteration] = fraction_accepted | |
| # Record attempt | |
| ncfile.variables['distance'][iteration] = final_distance / units.angstroms | |
| ncfile.variables['positions'][iteration,:,:] = coordinates[:,:] / units.angstroms | |
| ncfile.sync() | |
| # Debug. | |
| final_time = time.time() | |
| elapsed_time = final_time - initial_time | |
| print "%12.3f s elapsed" % elapsed_time | |
| # Increment iteration counter. | |
| iteration += 1 | |
| # Close netcdf file. | |
| ncfile.close() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/usr/local/bin/env python | |
| #============================================================================================= | |
| # MODULE DOCSTRING | |
| #============================================================================================= | |
| """ | |
| Sampling utility functions. | |
| DESCRIPTION | |
| COPYRIGHT | |
| @author John D. Chodera <jchodera@gmail.com> | |
| All code in this repository is released under the GNU General Public License. | |
| This program is free software: you can redistribute it and/or modify it under | |
| the terms of the GNU General Public License as published by the Free Software | |
| Foundation, either version 3 of the License, or (at your option) any later | |
| version. | |
| This program is distributed in the hope that it will be useful, but WITHOUT ANY | |
| WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A | |
| PARTICULAR PURPOSE. See the GNU General Public License for more details. | |
| You should have received a copy of the GNU General Public License along with | |
| this program. If not, see <http://www.gnu.org/licenses/>. | |
| TODO | |
| """ | |
| #============================================================================================= | |
| # GLOBAL IMPORTS | |
| #============================================================================================= | |
| import time | |
| import copy | |
| import numpy | |
| import simtk | |
| import simtk.unit as units | |
| import simtk.openmm as openmm | |
| #============================================================================================= | |
| # CONSTANTS | |
| #============================================================================================= | |
| kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA | |
| #============================================================================================= | |
| # Utility functions | |
| #============================================================================================= | |
| def compute_forces(context, positions): | |
| """ | |
| Compute forces for given positions. | |
| """ | |
| context.setPositions(positions) | |
| state = context.getState(getForces=True) | |
| forces = state.getForces(asNumpy=True) | |
| return forces | |
| def compute_energy(context, positions, velocities): | |
| """ | |
| Compute total energy for positions and velocities. | |
| """ | |
| # Set positions and velocities. | |
| context.setPositions(positions) | |
| context.setVelocities(velocities) | |
| # Compute total energy. | |
| state = context.getState(getEnergy=True) | |
| total_energy = state.getPotentialEnergy() + state.getKineticEnergy() | |
| #print "potential energy: %.3f kcal/mol" % (state.getPotentialEnergy() / units.kilocalories_per_mole) | |
| #print "kinetic energy: %.3f kcal/mol" % (state.getKineticEnergy() / units.kilocalories_per_mole) | |
| return total_energy | |
| def compute_forces_and_energy(context, positions, velocities): | |
| """ | |
| Compute total energy for positions and velocities. | |
| """ | |
| # Set positions and velocities. | |
| context.setPositions(positions) | |
| context.setVelocities(velocities) | |
| # Compute total energy. | |
| state = context.getState(getForces=True, getEnergy=True) | |
| forces = state.getForces(asNumpy=True) | |
| total_energy = state.getPotentialEnergy() + state.getKineticEnergy() | |
| #print "potential energy: %.3f kcal/mol" % (state.getPotentialEnergy() / units.kilocalories_per_mole) | |
| #print "kinetic energy: %.3f kcal/mol" % (state.getKineticEnergy() / units.kilocalories_per_mole) | |
| return [forces, total_energy] | |
| def compute_potential(context, positions): | |
| """ | |
| Compute potential energy for positions. | |
| """ | |
| # Set positions and velocities. | |
| context.setPositions(positions) | |
| # Compute total energy. | |
| state = context.getState(getEnergy=True) | |
| potential_energy = state.getPotentialEnergy() | |
| return potential_energy | |
| def minimize(platform, system, positions): | |
| # Create a Context. | |
| timestep = 1.0 * units.femtoseconds | |
| integrator = openmm.VerletIntegrator(timestep) | |
| context = openmm.Context(system, integrator, platform) | |
| # Set coordinates. | |
| context.setPositions(positions) | |
| # Compute initial energy. | |
| state = context.getState(getEnergy=True) | |
| initial_potential = state.getPotentialEnergy() | |
| print "initial potential: %12.3f kcal/mol" % (initial_potential / units.kilocalories_per_mole) | |
| # Minimize. | |
| openmm.LocalEnergyMinimizer.minimize(context) | |
| # Compute final energy. | |
| state = context.getState(getEnergy=True, getPositions=True) | |
| final_potential = state.getPotentialEnergy() | |
| positions = state.getPositions(asNumpy=True) | |
| # Report | |
| print "final potential : %12.3f kcal/mol" % (final_potential / units.kilocalories_per_mole) | |
| return positions | |
| #============================================================================================= | |
| # Equilibrate the system with OpenMM's leapfrog Langevin dynamics. | |
| #============================================================================================= | |
| def equilibrate_langevin(system, timestep, collision_rate, temperature, sqrt_kT_over_m, coordinates, platform): | |
| nsteps = 5000 | |
| print "Equilibrating for %.3f ps..." % ((nsteps * timestep) / units.picoseconds) | |
| # Create integrator and context. | |
| integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep) | |
| context = openmm.Context(system, integrator, platform) | |
| # Set coordinates. | |
| context.setPositions(coordinates) | |
| # Set Maxwell-Boltzmann velocities | |
| velocities = sqrt_kT_over_m * numpy.random.standard_normal(size=sqrt_kT_over_m.shape) | |
| context.setVelocities(velocities) | |
| # Equilibrate. | |
| integrator.step(nsteps) | |
| # Compute energy | |
| print "Computing energy." | |
| state = context.getState(getEnergy=True) | |
| potential_energy = state.getPotentialEnergy() | |
| print "potential energy: %.3f kcal/mol" % (potential_energy / units.kilocalories_per_mole) | |
| # Get coordinates. | |
| state = context.getState(getPositions=True, getVelocities=True) | |
| coordinates = state.getPositions(asNumpy=True) | |
| velocities = state.getVelocities(asNumpy=True) | |
| box_vectors = state.getPeriodicBoxVectors() | |
| system.setDefaultPeriodicBoxVectors(*box_vectors) | |
| print "Computing energy again." | |
| context.setPositions(coordinates) | |
| context.setVelocities(velocities) | |
| state = context.getState(getEnergy=True) | |
| potential_energy = state.getPotentialEnergy() | |
| print "potential energy: %.3f kcal/mol" % (potential_energy / units.kilocalories_per_mole) | |
| total_energy = compute_energy(context, coordinates, velocities) | |
| return [coordinates, velocities] | |
| #============================================================================================= | |
| # Equilibrate the system with hybrid Monte Carlo (HMC) dynamics. | |
| #============================================================================================= | |
| def equilibrate_hmc(system, timestep, collision_rate, temperature, masses, sqrt_kT_over_m, coordinates, platform, debug=False): | |
| nhmc = 100 # number of HMC iterations | |
| nsteps = 50 # number of steps per HMC iteration | |
| beta = 1.0 / (kB * temperature) | |
| print "Equilibrating for %d HMC moves of %.3f ps..." % (nhmc, (nsteps * timestep) / units.picoseconds) | |
| # Create integrator and context. | |
| integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep) | |
| context = openmm.Context(system, integrator, platform) | |
| naccepted = 0 | |
| for hmc_move in range(nhmc): | |
| if debug: print "HMC move %5d / %5d" % (hmc_move, nhmc) | |
| # Set coordinates. | |
| context.setPositions(coordinates) | |
| # Set Maxwell-Boltzmann velocities | |
| velocities = sqrt_kT_over_m * numpy.random.standard_normal(size=sqrt_kT_over_m.shape) | |
| context.setVelocities(velocities) | |
| # Compute initial total energy. | |
| state = context.getState(getEnergy=True) | |
| initial_energy = state.getPotentialEnergy() + state.getKineticEnergy() | |
| # Half-kick velocities backwards. | |
| state = context.getState(getForces=True) | |
| forces = state.getForces(asNumpy=True) | |
| velocities[:,:] -= 0.5 * forces[:,:]/masses[:,:] * timestep | |
| context.setVelocities(velocities) | |
| # Integrate using leapfrog. | |
| integrator.step(nsteps) | |
| # Half-kick velocities forwards. | |
| state = context.getState(getForces=True,getVelocities=True) | |
| forces = state.getForces(asNumpy=True) | |
| velocities = state.getVelocities(asNumpy=True) | |
| velocities[:,:] += 0.5 * forces[:,:]/masses[:,:] * timestep | |
| context.setVelocities(velocities) | |
| # Compute final energy. | |
| state = context.getState(getEnergy=True) | |
| final_energy = state.getPotentialEnergy() + state.getKineticEnergy() | |
| # Compute HMC acceptance. | |
| du = beta * (final_energy - initial_energy) | |
| if debug: print "du = %8.1f :" % du, | |
| if numpy.random.uniform() < numpy.exp(-du): | |
| # Accept. | |
| if debug: print " accepted." | |
| state = context.getState(getPositions=True) | |
| coordinates = state.getPositions(asNumpy=True) | |
| naccepted += 1 | |
| else: | |
| # Reject. | |
| if debug: print " rejected." | |
| pass | |
| # Compute energy | |
| if debug: print "Computing energy." | |
| state = context.getState(getEnergy=True) | |
| potential_energy = state.getPotentialEnergy() | |
| if debug: print "potential energy: %.3f kcal/mol" % (potential_energy / units.kilocalories_per_mole) | |
| # Get coordinates. | |
| state = context.getState(getPositions=True, getVelocities=True) | |
| coordinates = state.getPositions(asNumpy=True) | |
| velocities = state.getVelocities(asNumpy=True) | |
| box_vectors = state.getPeriodicBoxVectors() | |
| system.setDefaultPeriodicBoxVectors(*box_vectors) | |
| if debug: print "Computing energy again." | |
| context.setPositions(coordinates) | |
| context.setVelocities(velocities) | |
| state = context.getState(getEnergy=True) | |
| potential_energy = state.getPotentialEnergy() | |
| if debug: print "potential energy: %.3f kcal/mol" % (potential_energy / units.kilocalories_per_mole) | |
| total_energy = compute_energy(context, coordinates, velocities) | |
| fraction_accepted = float(naccepted) / float(nhmc) | |
| return [coordinates, velocities, fraction_accepted] | |
| #============================================================================================= | |
| # Generalized hybrid Monte Carlo (GHMC) integrator | |
| #============================================================================================= | |
| def equilibrate_ghmc(system, timestep, collision_rate, temperature, masses, sqrt_kT_over_m, positions, platform, debug=False): | |
| nsteps = 500 # number of steps | |
| kT = kB * temperature | |
| beta = 1.0 / kT # inverse temperature | |
| print "Equilibrating for %d GHMC steps (%.3f ps)..." % (nsteps, (nsteps * timestep) / units.picoseconds) | |
| initial_time = time.time() | |
| # Assign Maxwell-Boltzmann velocities | |
| velocities = sqrt_kT_over_m * numpy.random.standard_normal(size=sqrt_kT_over_m.shape) | |
| # Compute Langevin velocity modification factors. | |
| gamma = collision_rate * masses | |
| sigma2 = (2.0 * kT * gamma) | |
| alpha_factor = (1.0 - (timestep/4.0)*collision_rate) / (1.0 + (timestep/4.0)*collision_rate) | |
| x = (timestep/2.0*sigma2).in_units_of(units.kilogram**2/units.mole**2 * units.meter**2/units.second**2) | |
| y = units.Quantity(numpy.sqrt(x / x.unit), units.sqrt(1.0 * x.unit)) | |
| beta_factor = (y / (1.0 + (timestep/4.0)*collision_rate) / masses).in_units_of(velocities.unit) | |
| # Create integrator and context. | |
| integrator = openmm.VerletIntegrator(timestep) | |
| context = openmm.Context(system, integrator, platform) | |
| # Compute forces and total energy. | |
| context.setPositions(positions) | |
| context.setVelocities(velocities) | |
| state = context.getState(getForces=True, getEnergy=True) | |
| forces = state.getForces(asNumpy=True) | |
| kinetic_energy = state.getKineticEnergy() | |
| potential_energy = state.getPotentialEnergy() | |
| total_energy = kinetic_energy + potential_energy | |
| # Create storage for proposed positions and velocities. | |
| proposed_positions = copy.deepcopy(positions) | |
| proposed_velocities = copy.deepcopy(velocities) | |
| naccepted = 0 # number of accepted GHMC steps | |
| for step in range(nsteps): | |
| # | |
| # Velocity modification step. | |
| # | |
| velocities[:,:] = velocities[:,:]*alpha_factor + units.Quantity(numpy.random.standard_normal(size=positions.shape) * (beta_factor/beta_factor.unit), beta_factor.unit) | |
| kinetic_energy = 0.5 * (masses * velocities**2).in_units_of(potential_energy.unit).sum() * potential_energy.unit # have to do this because sum(...) and .sum() don't respect units | |
| total_energy = kinetic_energy + potential_energy | |
| # | |
| # Metropolis-wrapped Velocity Verlet step | |
| # | |
| proposed_positions[:,:] = positions[:,:] | |
| proposed_velocities[:,:] = velocities[:,:] | |
| # Half-kick velocities | |
| proposed_velocities[:,:] += 0.5 * forces[:,:]/masses[:,:] * timestep | |
| # Full-kick positions | |
| proposed_positions[:,:] += proposed_velocities[:,:] * timestep | |
| # Update force at new positions. | |
| context.setVelocities(proposed_velocities) | |
| context.setPositions(proposed_positions) | |
| state = context.getState(getForces=True, getEnergy=True) | |
| proposed_forces = state.getForces(asNumpy=True) | |
| proposed_potential_energy = state.getPotentialEnergy() | |
| # Half-kick velocities | |
| proposed_velocities[:,:] += 0.5 * proposed_forces[:,:]/masses[:,:] * timestep | |
| proposed_kinetic_energy = 0.5 * (masses * proposed_velocities**2).in_units_of(potential_energy.unit).sum() * potential_energy.unit # have to do this because sum(...) and .sum() don't respect units | |
| # Compute new total energy. | |
| proposed_total_energy = proposed_kinetic_energy + proposed_potential_energy | |
| # Accept or reject, inverting momentum if rejected. | |
| du = beta * (proposed_total_energy - total_energy) | |
| if (du < 0.0) or (numpy.random.uniform() < numpy.exp(-du)): | |
| # Accept and update positions, velocities, forces, and energies. | |
| naccepted += 1 | |
| positions[:,:] = proposed_positions[:,:] | |
| velocities[:,:] = proposed_velocities[:,:] | |
| forces[:,:] = proposed_forces[:,:] | |
| potential_energy = proposed_potential_energy | |
| kinetic_energy = proposed_kinetic_energy | |
| else: | |
| # Reject, requiring negation of velocities. | |
| velocities[:,:] = -velocities[:,:] | |
| # | |
| # Velocity modification step. | |
| # | |
| velocities[:,:] = velocities[:,:]*alpha_factor + units.Quantity(numpy.random.standard_normal(size=positions.shape) * (beta_factor/beta_factor.unit), beta_factor.unit) | |
| kinetic_energy = 0.5 * (masses * velocities**2).in_units_of(potential_energy.unit).sum() * potential_energy.unit # have to do this because sum(...) and .sum() don't respect units | |
| total_energy = kinetic_energy + potential_energy | |
| # Print final statistics. | |
| fraction_accepted = float(naccepted) / float(nsteps) | |
| final_time = time.time() | |
| elapsed_time = final_time - initial_time | |
| if debug: print "%12.3f s elapsed | accepted %6.3f%%" % (elapsed_time, fraction_accepted*100.0) | |
| return [positions, velocities, fraction_accepted] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/usr/local/bin/env python | |
| #============================================================================================= | |
| # MODULE DOCSTRING | |
| #============================================================================================= | |
| """ | |
| WCA fluid and WCA dimer systems. | |
| DESCRIPTION | |
| COPYRIGHT | |
| @author John D. Chodera <jchodera@gmail.com> | |
| All code in this repository is released under the GNU General Public License. | |
| This program is free software: you can redistribute it and/or modify it under | |
| the terms of the GNU General Public License as published by the Free Software | |
| Foundation, either version 3 of the License, or (at your option) any later | |
| version. | |
| This program is distributed in the hope that it will be useful, but WITHOUT ANY | |
| WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A | |
| PARTICULAR PURPOSE. See the GNU General Public License for more details. | |
| You should have received a copy of the GNU General Public License along with | |
| this program. If not, see <http://www.gnu.org/licenses/>. | |
| TODO | |
| """ | |
| #============================================================================================= | |
| # GLOBAL IMPORTS | |
| #============================================================================================= | |
| import numpy | |
| import simtk | |
| import simtk.unit as units | |
| import simtk.openmm as openmm | |
| #============================================================================================= | |
| # CONSTANTS | |
| #============================================================================================= | |
| kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA | |
| #============================================================================================= | |
| # DEFAULT PARAMETERS | |
| #============================================================================================= | |
| natoms = 216 # number of particles | |
| # WCA fluid parameters (argon). | |
| mass = 39.9 * units.amu # reference mass | |
| sigma = 3.4 * units.angstrom # reference lengthscale | |
| epsilon = 120.0 * units.kelvin * kB # reference energy | |
| r_WCA = 2.**(1./6.) * sigma # interaction truncation range | |
| tau = units.sqrt(sigma**2 * mass / epsilon) # characteristic timescale | |
| # Simulation conditions. | |
| temperature = 0.824 / (kB / epsilon) # temperature | |
| kT = kB * temperature # thermal energy | |
| beta = 1.0 / kT # inverse temperature | |
| density = 0.96 / sigma**3 # default density | |
| stable_timestep = 0.001 * tau # stable timestep | |
| collision_rate = 1 / tau # collision rate for Langevin interator | |
| # Dimer potential parameters. | |
| h = 9.0 * kT # barrier height | |
| #h = 5.0 * kT # barrier height (DEBUG) | |
| r0 = r_WCA # compact state distance | |
| w = 0.5 * r_WCA # extended state distance is r0 + 2*w | |
| #============================================================================================= | |
| # WCA Fluid | |
| #============================================================================================= | |
| def WCAFluid(N=natoms, density=density, mm=None, mass=mass, epsilon=epsilon, sigma=sigma): | |
| """ | |
| Create a Weeks-Chandler-Andersen system. | |
| OPTIONAL ARGUMENTS | |
| N (int) - total number of atoms (default: 150) | |
| density (float) - N sigma^3 / V (default: 0.96) | |
| sigma | |
| epsilon | |
| """ | |
| # Choose OpenMM package. | |
| if mm is None: | |
| mm = openmm | |
| # Create system | |
| system = mm.System() | |
| # Compute total system volume. | |
| volume = N / density | |
| # Make system cubic in dimension. | |
| length = volume**(1.0/3.0) | |
| # TODO: Can we change this to use tuples or 3x3 array? | |
| a = units.Quantity(numpy.array([1.0, 0.0, 0.0], numpy.float32), units.nanometer) * length/units.nanometer | |
| b = units.Quantity(numpy.array([0.0, 1.0, 0.0], numpy.float32), units.nanometer) * length/units.nanometer | |
| c = units.Quantity(numpy.array([0.0, 0.0, 1.0], numpy.float32), units.nanometer) * length/units.nanometer | |
| print "box edge length = %s" % str(length) | |
| system.setDefaultPeriodicBoxVectors(a, b, c) | |
| # Add particles to system. | |
| for n in range(N): | |
| system.addParticle(mass) | |
| # Create nonbonded force term implementing Kob-Andersen two-component Lennard-Jones interaction. | |
| energy_expression = '4.0*epsilon*((sigma/r)^12 - (sigma/r)^6) + epsilon' | |
| # Create force. | |
| force = mm.CustomNonbondedForce(energy_expression) | |
| # Set epsilon and sigma global parameters. | |
| force.addGlobalParameter('epsilon', epsilon) | |
| force.addGlobalParameter('sigma', sigma) | |
| # Add particles | |
| for n in range(N): | |
| force.addParticle([]) | |
| # Set periodic boundary conditions with cutoff. | |
| force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffNonPeriodic) | |
| print "setting cutoff distance to %s" % str(r_WCA) | |
| force.setCutoffDistance(r_WCA) | |
| # Add nonbonded force term to the system. | |
| system.addForce(force) | |
| # Create initial coordinates using random positions. | |
| coordinates = units.Quantity(numpy.random.rand(N,3), units.nanometer) * (length / units.nanometer) | |
| # Return system and coordinates. | |
| return [system, coordinates] | |
| #============================================================================================= | |
| # WCA dimer plus fluid | |
| #============================================================================================= | |
| def WCADimer(N=natoms, density=density, mm=None, mass=mass, epsilon=epsilon, sigma=sigma, h=h, r0=r0, w=w): | |
| """ | |
| Create a bistable bonded pair of particles (indices 0 and 1) optionally surrounded by a Weeks-Chandler-Andersen fluid. | |
| The bistable potential has form | |
| U(r) = h*(1-((r-r0-w)/w)^2)^2 | |
| where r0 is the compact state separation, r0+2w is the extended state separation, and h is the barrier height. | |
| The WCA potential has form | |
| U(r) = 4 epsilon [ (sigma/r)^12 - (sigma/r)^6 ] + epsilon (r < r*) | |
| = 0 (r >= r*) | |
| where r* = 2^(1/6) sigma. | |
| OPTIONAL ARGUMENTS | |
| N (int) - total number of atoms (default: 2) | |
| density (float) - number density of particles (default: 0.96 / sigma**3) | |
| mass (simtk.unit.Quantity of mass) - particle mass (default: 39.948 amu) | |
| sigma (simtk.unit.Quantity of length) - Lennard-Jones sigma parameter (default: 0.3405 nm) | |
| epsilon (simtk.unit.Quantity of energy) - Lennard-Jones well depth (default: (119.8 Kelvin)*kB) | |
| h (simtk.unit.Quantity of energy) - bistable potential barrier height (default: ???) | |
| r0 (simtk.unit.Quantity of length) - bistable potential compact state separation (default: ???) | |
| w (simtk.unit.Quantity of length) - bistable potential extended state separation is r0+2*w (default: ???) | |
| """ | |
| # Choose OpenMM package. | |
| if mm is None: | |
| mm = openmm | |
| # Compute cutoff for WCA fluid. | |
| r_WCA = 2.**(1./6.) * sigma # cutoff at minimum of potential | |
| # Create system | |
| system = mm.System() | |
| # Compute total system volume. | |
| volume = N / density | |
| # Make system cubic in dimension. | |
| length = volume**(1.0/3.0) | |
| a = units.Quantity(numpy.array([1.0, 0.0, 0.0], numpy.float32), units.nanometer) * length/units.nanometer | |
| b = units.Quantity(numpy.array([0.0, 1.0, 0.0], numpy.float32), units.nanometer) * length/units.nanometer | |
| c = units.Quantity(numpy.array([0.0, 0.0, 1.0], numpy.float32), units.nanometer) * length/units.nanometer | |
| print "box edge length = %s" % str(length) | |
| system.setDefaultPeriodicBoxVectors(a, b, c) | |
| # Add particles to system. | |
| for n in range(N): | |
| system.addParticle(mass) | |
| # WCA: Lennard-Jones truncated at minim and shifted so potential is zero at cutoff. | |
| energy_expression = '4.0*epsilon*((sigma/r)^12 - (sigma/r)^6) + epsilon' | |
| # Create force. | |
| force = mm.CustomNonbondedForce(energy_expression) | |
| # Set epsilon and sigma global parameters. | |
| force.addGlobalParameter('epsilon', epsilon) | |
| force.addGlobalParameter('sigma', sigma) | |
| # Add particles | |
| for n in range(N): | |
| force.addParticle([]) | |
| # Add exclusion between bonded particles. | |
| force.addExclusion(0,1) | |
| # Set periodic boundary conditions with cutoff. | |
| if (N > 2): | |
| force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic) | |
| else: | |
| force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffNonPeriodic) | |
| print "setting cutoff distance to %s" % str(r_WCA) | |
| force.setCutoffDistance(r_WCA) | |
| # Add nonbonded force term to the system. | |
| system.addForce(force) | |
| # Add dimer potential to first two particles. | |
| dimer_force = openmm.CustomBondForce('h*(1-((r-r0-w)/w)^2)^2;') | |
| dimer_force.addGlobalParameter('h', h) # barrier height | |
| dimer_force.addGlobalParameter('r0', r0) # compact state separation | |
| dimer_force.addGlobalParameter('w', w) # second minimum is at r0 + 2*w | |
| dimer_force.addBond(0, 1, []) | |
| system.addForce(dimer_force) | |
| # Create initial coordinates using random positions. | |
| coordinates = units.Quantity(numpy.random.rand(N,3), units.nanometer) * (length / units.nanometer) | |
| # Reposition dimer particles at compact minimum. | |
| coordinates[0,:] *= 0.0 | |
| coordinates[1,:] *= 0.0 | |
| coordinates[1,0] = r0 | |
| # Return system and coordinates. | |
| return [system, coordinates] | |
| #============================================================================================= | |
| # WCA dimer in vacuum | |
| #============================================================================================= | |
| def WCADimerVacuum(mm=None, mass=mass, epsilon=epsilon, sigma=sigma, h=h, r0=r0, w=w): | |
| """ | |
| Create a bistable dimer. | |
| OPTIONAL ARGUMENTS | |
| """ | |
| # Choose OpenMM package. | |
| if mm is None: | |
| mm = openmm | |
| # Create system | |
| system = mm.System() | |
| # Add particles to system. | |
| for n in range(2): | |
| system.addParticle(mass) | |
| # Add dimer potential to first two particles. | |
| dimer_force = openmm.CustomBondForce('h*(1-((r-r0-w)/w)^2)^2;') | |
| dimer_force.addGlobalParameter('h', h) # barrier height | |
| dimer_force.addGlobalParameter('r0', r0) # compact state separation | |
| dimer_force.addGlobalParameter('w', w) # second minimum is at r0 + 2*w | |
| dimer_force.addBond(0, 1, []) | |
| system.addForce(dimer_force) | |
| # Create initial coordinates using random positions. | |
| coordinates = units.Quantity(numpy.zeros([2,3], numpy.float64), units.nanometer) | |
| # Reposition dimer particles at compact minimum. | |
| coordinates[0,:] *= 0.0 | |
| coordinates[1,:] *= 0.0 | |
| coordinates[1,0] = r0 | |
| # Return system and coordinates. | |
| return [system, coordinates] | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment