#!/usr/local/bin/env python
#=============================================================================================
# MODULE DOCSTRING
#=============================================================================================
"""
Compare platforms using different tolerances.
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 simtk.unit as units
import simtk.chem.openmm as openmm
#import simtk.chem.openmm.extras.amber as amber
import amber
#=============================================================================================
# SUBROUTINES
#=============================================================================================
def setIonCharge(system, atom_index, new_charge):
"""
Set ion charge to specified value.
"""
for force_index in range(system.getNumForces()):
force = system.getForce(force_index)
if hasattr(force, 'getParticleParameters'):
[old_charge, sigma, epsilon] = force.getParticleParameters(atom_index)
force.setParticleParameters(atom_index, new_charge, sigma, epsilon)
break
return
def equilibrate(system, coordinates, platform):
print "Equilibrating..."
# Create integrator and context.
integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
context = openmm.Context(system, integrator, platform)
# Set coordinates.
context.setPositions(coordinates)
# Equilibrate.
integrator.step(5000) # 10 ps
# 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)
return [coordinates, velocities]
#=============================================================================================
# MAIN AND TESTS
#=============================================================================================
if __name__ == "__main__":
# Parameters
prmtop_filename = 'system.prmtop'
crd_filename = 'system.crd'
temperature = 298.0 * units.kelvin
collision_rate = 5.0 / units.picoseconds
timestep = 2.0 * units.femtoseconds
# 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("Cuda")
# Create system.
print "Reading system..."
system = amber.readAmberSystem(prmtop_filename, nonbondedMethod='reaction-field', nonbondedCutoff=9.0*units.angstroms, shake='h-bonds')
[coordinates, box_vectors] = amber.readAmberCoordinates(crd_filename, read_box=True)
system.setPeriodicBoxVectors(box_vectors[0], box_vectors[1], box_vectors[2])
# Equilibrate.
setIonCharge(system, 0, -1.0 * units.elementary_charge)
for iteration in range(100):
[coordinates, velocities] = equilibrate(system, coordinates, platform)
# Generate work samples.
outfile = open('work.reverse', 'w') # file for writing work values
lechner_outfile = open('work.reverse.lechner', 'w') # file for writing work values
niterations = 1000 # number of work samples to collect
nsteps_to_try = [1, 5, 10, 50, 100, 500, 1000, 5000] # number of steps per switch
for iteration in range(niterations):
print "iteration %5d / %5d" % (iteration, niterations)
# Equilibrate.
setIonCharge(system, 0, -1.0 * units.elementary_charge)
[coordinates, velocities] = equilibrate(system, coordinates, platform)
for nsteps in nsteps_to_try:
# Accumulate switching work.
work = 0.0 * units.kilocalories_per_mole
#integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
integrator = openmm.VerletIntegrator(timestep)
proposed_coordinates = copy.deepcopy(coordinates)
proposed_velocities = copy.deepcopy(velocities)
# Accumulate Lechner work
setIonCharge(system, 0, -1.0 * units.elementary_charge)
context = openmm.Context(system, integrator, platform)
context.setPositions(proposed_coordinates)
context.setVelocities(proposed_velocities)
state = context.getState(getEnergy=True)
work_lechner = - (state.getPotentialEnergy() + state.getKineticEnergy())
del state, context
for step in range(nsteps):
# Compute old energy.
setIonCharge(system, 0, (-1.0 + 2.0*float(step)/float(nsteps)) * units.elementary_charge)
context = openmm.Context(system, integrator, platform)
context.setPositions(proposed_coordinates)
context.setVelocities(proposed_velocities)
state = context.getState(getEnergy=True)
old_energy = state.getPotentialEnergy() + state.getKineticEnergy()
del state, context
# Compute new energy.
setIonCharge(system, 0, (-1.0 + 2.0*float(step+1)/float(nsteps)) * units.elementary_charge)
context = openmm.Context(system, integrator, platform)
context.setPositions(proposed_coordinates)
context.setVelocities(proposed_velocities)
state = context.getState(getEnergy=True)
new_energy = state.getPotentialEnergy() + state.getKineticEnergy()
#del state, context
# Evolve dynamics.
#context = openmm.Context(system, integrator, platform)
#context.setPositions(coordinates)
integrator.step(1)
state = context.getState(getPositions=True, getVelocities=True)
proposed_coordinates = state.getPositions(asNumpy=True)
proposed_velocities = state.getVelocities(asNumpy=True)
del state, context
# Accumulate work.
delta_energy = new_energy - old_energy
work += delta_energy
#print "step %5d / %5d : energy = %8.1f kcal/mol, delta_energy = %8.1f kcal/mol, accumulated work = %8.1f kcal/mol" % (step+1, nsteps, new_energy / units.kilocalories_per_mole, delta_energy / units.kilocalories_per_mole, work / units.kilocalories_per_mole)
outfile.write("%12.3f" % (work / units.kilocalories_per_mole))
# Accumulate Lechner work
setIonCharge(system, 0, 1.0 * units.elementary_charge)
context = openmm.Context(system, integrator, platform)
context.setPositions(proposed_coordinates)
context.setVelocities(proposed_velocities)
state = context.getState(getEnergy=True)
work_lechner += state.getPotentialEnergy() + state.getKineticEnergy()
del state, context
lechner_outfile.write("%12.3f" % (work_lechner / units.kilocalories_per_mole))
outfile.write("\n")
outfile.flush()
lechner_outfile.write("\n")
lechner_outfile.flush()
outfile.close()
lechner_outfile.close()