<html><head><meta name="color-scheme" content="light dark"></head><body><pre style="word-wrap: break-word; white-space: pre-wrap;">#!/usr/local/bin/env python

#=============================================================================================
# MODULE DOCSTRING
#=============================================================================================

"""
Compare platforms using different tolerances.

DESCRIPTION

COPYRIGHT

@author John D. Chodera &lt;jchodera@gmail.com&gt;

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 &lt;http://www.gnu.org/licenses/&gt;.

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(10000) # 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
    pressure = 1.0 * units.atmospheres
    collision_rate = 5.0 / units.picoseconds
    timestep = 1.0 * units.femtoseconds
    barostat_frequency = 25

    # 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")

    # Equilibrate with a barostat to determine appropriate box size.
    print "Creating and equilibrating system with barostat.."
    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.setDefaultPeriodicBoxVectors(box_vectors[0], box_vectors[1], box_vectors[2])
    # Create barostat.
    barostat = openmm.MonteCarloBarostat(pressure, temperature, barostat_frequency)
    system.addForce(barostat)
    # Equilibrate.
    setIonCharge(system, 0, 1.0 * units.elementary_charge)
    integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    context = openmm.Context(system, integrator, platform)
    context.setPositions(coordinates)
    integrator.step(10000) # 10 ps
    state = context.getState(getPositions=True)
    coordinates = state.getPositions(asNumpy=True)
    box_vectors = state.getPeriodicBoxVectors()
    del state, context, integrator
    
    # Create constant-volume system.
    system = amber.readAmberSystem(prmtop_filename, nonbondedMethod='reaction-field', nonbondedCutoff=9.0*units.angstroms, shake='h-bonds')
    system.setDefaultPeriodicBoxVectors(box_vectors[0], box_vectors[1], box_vectors[2])    

    # Equilibrate NVT.
    setIonCharge(system, 0, 1.0 * units.elementary_charge)
    for iteration in range(10):
        [coordinates, velocities] = equilibrate(system, coordinates, platform)

    # Generate work samples.
    outfile = open('work.forward', 'w') # file for writing work values
    lechner_outfile = open('work.forward.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)

        print coordinates

        for nsteps in nsteps_to_try:
            # Accumulate switching work.
            work = 0.0 * units.kilocalories_per_mole
            heat = 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()
                
                # Evolve dynamics.
                integrator.step(1)
                state = context.getState(getPositions=True, getVelocities=True)    
                proposed_coordinates = state.getPositions(asNumpy=True)
                proposed_velocities = state.getVelocities(asNumpy=True)                                

                # Compute energy after dynamics.
                state = context.getState(getEnergy=True)
                energy_after_dynamics = state.getPotentialEnergy() + state.getKineticEnergy()
                heat += (energy_after_dynamics - new_energy)                

                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))

            print "%d steps : work %12.3f kcal/mol, heat %12.3f kcal/mol, total %12.3f kcal/mol" % (nsteps, work / units.kilocalories_per_mole, heat / units.kilocalories_per_mole, work_lechner / units.kilocalories_per_mole)

        outfile.write("\n")
        outfile.flush()
        lechner_outfile.write("\n")
        lechner_outfile.flush()
        
    outfile.close()
    lechner_outfile.close()
</pre></body></html>